diff --git a/.env b/.env index d3e1c1d63a..ef0c5fb101 100644 --- a/.env +++ b/.env @@ -40,18 +40,17 @@ ARCH_SHORT=amd64 # Default repository to pull and push images from REPO=ghcr.io/apache/arrow-java-dev -ARROW_REPO=apache/arrow-dev +ARROW_REPO=ghcr.io/apache/arrow-dev # The setup attempts to generate coredumps by default, in order to disable the # coredump generation set it to 0 ULIMIT_CORE=-1 # Default versions for various dependencies -JDK=11 +JDK=17 MAVEN=3.9.9 # Versions for various dependencies used to build artifacts # Keep in sync with apache/arrow ARROW_REPO_ROOT=./arrow -PYTHON=3.9 -VCPKG="f7423ee180c4b7f40d43402c2feb3859161ef625" # 2024.06.15 Release +VCPKG="9b965a116838c6cdcd36bca60d1b81b030c8ab8d" # 2026.05.27 (not release, upstream commit) diff --git a/.github/workflows/comment_bot.yml b/.github/workflows/comment_bot.yml index 5fbc858cc6..507d6a969c 100644 --- a/.github/workflows/comment_bot.yml +++ b/.github/workflows/comment_bot.yml @@ -30,7 +30,7 @@ jobs: if: github.event.comment.body == 'take' runs-on: ubuntu-latest steps: - - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + - uses: actions/github-script@v9 with: github-token: ${{ secrets.GITHUB_TOKEN }} script: |- diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 37320898a6..f4dc9ad6f1 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -19,7 +19,9 @@ name: Dev on: pull_request: {} - push: {} + push: + branches-ignore: + - dependabot/** concurrency: group: ${{ github.repository }}-${{ github.ref }}-${{ github.workflow }} @@ -33,16 +35,16 @@ jobs: name: "pre-commit" runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 with: fetch-depth: 0 persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v7 with: python-version: '3.x' - name: pre-commit (cache) - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ~/.cache/pre-commit key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }} diff --git a/.github/workflows/dev_pr.yml b/.github/workflows/dev_pr.yml index 34b3363c50..ad000df88e 100644 --- a/.github/workflows/dev_pr.yml +++ b/.github/workflows/dev_pr.yml @@ -35,6 +35,7 @@ concurrency: permissions: contents: read + issues: write pull-requests: write jobs: @@ -42,35 +43,35 @@ jobs: name: "Ensure PR format" runs-on: ubuntu-latest steps: - - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - uses: actions/checkout@v7 with: fetch-depth: 0 persist-credentials: false - name: Ensure PR title format id: title-format - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + uses: actions/github-script@v9 with: script: | const scripts = require(`${process.env.GITHUB_WORKSPACE}/.github/workflows/dev_pr.js`); return scripts.check_title_format({core, github, context}); - name: Label PR - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + uses: actions/github-script@v9 with: script: | const scripts = require(`${process.env.GITHUB_WORKSPACE}/.github/workflows/dev_pr.js`); await scripts.apply_labels({core, github, context}); - name: Ensure PR is labeled - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + uses: actions/github-script@v9 with: script: | const scripts = require(`${process.env.GITHUB_WORKSPACE}/.github/workflows/dev_pr.js`); await scripts.check_labels({core, github, context}); - name: Ensure PR is linked to an issue - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + uses: actions/github-script@v9 with: script: | const scripts = require(`${process.env.GITHUB_WORKSPACE}/.github/workflows/dev_pr.js`); diff --git a/.github/workflows/dev_pr_milestone.sh b/.github/workflows/dev_pr_milestone.sh index b6876b4b08..4a77eb1f73 100755 --- a/.github/workflows/dev_pr_milestone.sh +++ b/.github/workflows/dev_pr_milestone.sh @@ -37,8 +37,8 @@ main() { local -r milestone=$( gh api "/repos/${repo}/milestones" | jq --raw-output '.[] | .title' | - grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' - head -n1 + grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' | + head -n1 ) echo "Assigning milestone: ${milestone}" diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml new file mode 100644 index 0000000000..3872d03d2f --- /dev/null +++ b/.github/workflows/integration.yml @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +name: Integration + +on: + push: + branches: + - '**' + - '!dependabot/**' + tags: + - '**' + paths: + - '.github/workflows/integration.yml' + - '**/pom.xml' + - 'c/**' + - 'ci/scripts/**' + - 'compose.yaml' + - 'flight/**' + - 'format/**' + - 'testing/data/**' + - 'vector/**' + pull_request: + paths: + - '.github/workflows/integration.yml' + - '**/pom.xml' + - 'c/**' + - 'ci/scripts/**' + - 'compose.yaml' + - 'flight/**' + - 'format/**' + - 'testing/data/**' + - 'vector/**' + +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +permissions: + contents: read + +env: + DOCKER_VOLUME_PREFIX: ".docker/" + +jobs: + integration: + name: AMD64 integration + runs-on: ubuntu-latest + timeout-minutes: 60 + steps: + - name: Checkout Arrow + uses: actions/checkout@v7 + with: + fetch-depth: 0 + repository: apache/arrow + submodules: recursive + - name: Checkout Arrow Rust + uses: actions/checkout@v7 + with: + repository: apache/arrow-rs + path: rust + - name: Checkout Arrow nanoarrow + uses: actions/checkout@v7 + with: + repository: apache/arrow-nanoarrow + path: nanoarrow + - name: Checkout Arrow .NET + uses: actions/checkout@v7 + with: + repository: apache/arrow-dotnet + path: dotnet + - name: Checkout Arrow Go + uses: actions/checkout@v7 + with: + repository: apache/arrow-go + path: go + - name: Checkout Arrow Java + uses: actions/checkout@v7 + with: + path: java + - name: Checkout Arrow JavaScript + uses: actions/checkout@v7 + with: + repository: apache/arrow-js + path: js + - name: Free up disk space + run: | + ci/scripts/util_free_space.sh + - name: Cache Docker Volumes + uses: actions/cache@v6 + with: + path: .docker + key: integration-conda-${{ hashFiles('cpp/**') }} + restore-keys: integration-conda- + - name: Setup Python + uses: actions/setup-python@v6 + with: + python-version: 3.12 + - name: Setup Archery + run: pip install -e dev/archery[docker] + - name: Execute Docker Build + run: | + source ci/scripts/util_enable_core_dumps.sh + archery docker run \ + -e ARCHERY_DEFAULT_BRANCH=main \ + -e ARCHERY_INTEGRATION_TARGET_IMPLEMENTATIONS=java \ + -e ARCHERY_INTEGRATION_WITH_DOTNET=1 \ + -e ARCHERY_INTEGRATION_WITH_GO=1 \ + -e ARCHERY_INTEGRATION_WITH_JAVA=1 \ + -e ARCHERY_INTEGRATION_WITH_JS=1 \ + -e ARCHERY_INTEGRATION_WITH_NANOARROW=1 \ + -e ARCHERY_INTEGRATION_WITH_RUST=1 \ + conda-integration diff --git a/.github/workflows/rc.yml b/.github/workflows/rc.yml index 7e3cf5f6f2..18a721ac02 100644 --- a/.github/workflows/rc.yml +++ b/.github/workflows/rc.yml @@ -38,7 +38,7 @@ jobs: timeout-minutes: 5 steps: - name: Checkout - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: submodules: recursive - name: Prepare for tag @@ -71,7 +71,7 @@ jobs: run: | dev/release/run_rat.sh "${TAR_GZ}" - name: Upload source archive - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: release-source path: | @@ -101,45 +101,39 @@ jobs: packages: write steps: - name: Download source archive - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: release-source - name: Extract source archive run: | tar -xf apache-arrow-java-*.tar.gz --strip-components=1 - # We always use the main branch for apache/arrow for now. - # Because we want to use - # https://github.com/apache/arrow/pull/45114 in - # apache/arrow-java. We can revert this workaround once Apache - # Arrow 20.0.0 that includes the change released. - # - # - name: Download the latest Apache Arrow C++ - # if: github.event_name != 'schedule' - # run: | - # ci/scripts/download_cpp.sh + - name: Download the latest Apache Arrow C++ + if: github.event_name != 'schedule' + run: | + ci/scripts/download_cpp.sh - name: Checkout Apache Arrow C++ - # if: github.event_name == 'schedule' - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + if: github.event_name == 'schedule' + uses: actions/checkout@v7 with: repository: apache/arrow path: arrow - name: Checkout apache/arrow-testing - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/arrow-testing path: arrow/testing - name: Checkout apache/parquet-testing - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/parquet-testing path: arrow/cpp/submodules/parquet-testing - - uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 # v3.4.0 + - uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0 with: registry: ghcr.io username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - name: Cache - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: .docker key: jni-linux-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} @@ -154,54 +148,54 @@ jobs: - name: Compress into single artifact to keep directory structure run: tar -cvzf jni-linux-${{ matrix.platform.arch }}.tar.gz jni/ - name: Upload artifacts - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: jni-linux-${{ matrix.platform.arch }} path: jni-linux-${{ matrix.platform.arch }}.tar.gz jni-macos: name: JNI ${{ matrix.platform.runs_on }} ${{ matrix.platform.arch }} runs-on: ${{ matrix.platform.runs_on }} - timeout-minutes: 45 + timeout-minutes: 60 needs: - source strategy: fail-fast: false matrix: platform: - - { runs_on: macos-13, arch: "x86_64"} + - { runs_on: macos-15-intel, arch: "x86_64"} - { runs_on: macos-14, arch: "aarch_64" } env: MACOSX_DEPLOYMENT_TARGET: "14.0" steps: - name: Download source archive - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: release-source - name: Extract source archive run: | tar -xf apache-arrow-java-*.tar.gz --strip-components=1 - # - name: Download the latest Apache Arrow C++ - # if: github.event_name != 'schedule' - # run: | - # ci/scripts/download_cpp.sh + - name: Download the latest Apache Arrow C++ + if: github.event_name != 'schedule' + run: | + ci/scripts/download_cpp.sh - name: Checkout Apache Arrow C++ - # if: github.event_name == 'schedule' - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + if: github.event_name == 'schedule' + uses: actions/checkout@v7 with: repository: apache/arrow path: arrow - name: Checkout apache/arrow-testing - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/arrow-testing path: arrow/testing - name: Checkout apache/parquet-testing - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/parquet-testing path: arrow/cpp/submodules/parquet-testing - name: Set up Python - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + uses: actions/setup-python@v7 with: cache: 'pip' python-version: 3.12 @@ -221,6 +215,10 @@ jobs: # llvm@14 because llvm is newer than llvm@14. brew uninstall llvm || : + # We can remove this when we drop support for + # macos-15-intel. because macos-14 or later with arm64 uses /opt/homebrew/ + # not /usr/local/. + # # Ensure updating python@XXX with the "--overwrite" option. # If python@XXX is updated without "--overwrite", it causes # a conflict error. Because Python 3 installed not by @@ -229,10 +227,10 @@ jobs: # tries to replace /usr/local/bin/2to3 and so on and causes # a conflict error. brew update - for python_package in $(brew list | grep python@); do + for python_package in $(brew list | grep python@ | sort -r); do brew install --overwrite ${python_package} done - brew install --overwrite python + brew install --overwrite python3 if [ "$(uname -m)" = "arm64" ]; then # pkg-config formula is deprecated but it's still installed @@ -266,7 +264,7 @@ jobs: run: | echo "CCACHE_DIR=${PWD}/ccache" >> ${GITHUB_ENV} - name: Cache ccache - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ccache key: jni-macos-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} @@ -275,12 +273,12 @@ jobs: run: | set -e # make brew Java available to CMake - export JAVA_HOME=$(brew --prefix openjdk@11)/libexec/openjdk.jdk/Contents/Home + export JAVA_HOME=$(brew --prefix openjdk@17)/libexec/openjdk.jdk/Contents/Home ci/scripts/jni_macos_build.sh . arrow build jni - name: Compress into single artifact to keep directory structure run: tar -cvzf jni-macos-${{ matrix.platform.arch }}.tar.gz jni/ - name: Upload artifacts - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: jni-macos-${{ matrix.platform.arch }} path: jni-macos-${{ matrix.platform.arch }}.tar.gz @@ -294,11 +292,11 @@ jobs: fail-fast: false matrix: platform: - - runs_on: windows-2019 + - runs_on: windows-2022 arch: "x86_64" steps: - name: Download source archive - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: release-source - name: Extract source archive @@ -312,14 +310,14 @@ jobs: ci/scripts/download_cpp.sh - name: Checkout Apache Arrow C++ if: github.event_name == 'schedule' - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/arrow path: arrow - name: Set up Java - uses: actions/setup-java@7a6d8a8234af8eb26422e24e3006232cccaa061b # v4.6.0 + uses: actions/setup-java@v5 with: - java-version: '11' + java-version: '17' distribution: 'temurin' - name: Download Timezone Database shell: bash @@ -342,7 +340,7 @@ jobs: run: | echo "CCACHE_DIR=${PWD}/ccache" >> ${GITHUB_ENV} - name: Cache ccache - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ccache key: jni-windows-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} @@ -350,7 +348,7 @@ jobs: - name: Build shell: cmd run: | - call "C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" x64 + call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" x64 REM For ORC set TZDIR=/c/msys64/usr/share/zoneinfo bash -c "ci/scripts/jni_windows_build.sh . arrow build jni" @@ -358,7 +356,7 @@ jobs: shell: bash run: tar -cvzf jni-windows-${{ matrix.platform.arch }}.tar.gz jni/ - name: Upload artifacts - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: jni-windows-${{ matrix.platform.arch }} path: jni-windows-${{ matrix.platform.arch }}.tar.gz @@ -371,7 +369,7 @@ jobs: - jni-windows steps: - name: Download artifacts - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: path: artifacts - name: Decompress artifacts @@ -411,12 +409,12 @@ jobs: test -f jni/arrow_dataset_jni/x86_64/arrow_dataset_jni.dll test -f jni/arrow_orc_jni/x86_64/arrow_orc_jni.dll - name: Checkout apache/arrow-testing - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: repository: apache/arrow-testing path: testing - name: Cache ~/.m2 - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ~/.m2 key: binaries-build-${{ hashFiles('**/*.java', '**/pom.xml') }} @@ -430,12 +428,12 @@ jobs: cp -a target/site/apidocs reference tar -cvzf reference.tar.gz reference - name: Upload binaries - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: release-binaries path: binaries/* - name: Upload docs - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: reference path: reference.tar.gz @@ -448,15 +446,15 @@ jobs: contents: read packages: write steps: - - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 + - uses: actions/setup-python@v7 with: cache: 'pip' - name: Download source archive - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: release-source - name: Download Javadocs - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: reference - name: Extract source archive @@ -473,7 +471,7 @@ jobs: - name: Compress into single artifact to keep directory structure run: tar -cvzf html.tar.gz -C docs/build html - name: Upload artifacts - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: release-html path: html.tar.gz @@ -499,7 +497,7 @@ jobs: contents: write steps: - name: Checkout - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: path: site - name: Prepare branch @@ -521,7 +519,7 @@ jobs: cp ../.asf.yaml ./ git add .nojekyll .asf.yaml - name: Download - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: release-html - name: Extract @@ -557,7 +555,7 @@ jobs: - ubuntu-latest steps: - name: Download release artifacts - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: pattern: release-* - name: Verify @@ -591,7 +589,7 @@ jobs: contents: write steps: - name: Download release artifacts - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: pattern: release-* path: artifacts diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a955964cd8..7692eb6cbe 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -32,7 +32,7 @@ jobs: publish: name: Publish runs-on: ubuntu-latest - timeout-minutes: 5 + timeout-minutes: 30 steps: - name: Download RC contents run: | @@ -65,7 +65,7 @@ jobs: $artifact done - name: Checkout for publishing docs - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + uses: actions/checkout@v7 with: path: site - name: Publish docs diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5db5c988eb..653b16fa32 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -45,7 +45,7 @@ jobs: strategy: fail-fast: false matrix: - jdk: [11, 17, 21, 23] + jdk: [17, 21, 23] maven: [3.9.9] image: [ubuntu, conda-jni-cdata] include: @@ -58,15 +58,15 @@ jobs: MAVEN: ${{ matrix.maven }} steps: - name: Checkout Arrow - uses: actions/checkout@eef61447b9ff4aafe5dcd4e0bbf5d482be7e7871 # v4.2.1 + uses: actions/checkout@v7 with: fetch-depth: 0 submodules: recursive - name: Cache Docker Volumes - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 + uses: actions/cache@v6 with: path: .docker - key: maven-${{ matrix.jdk }}-${{ matrix.maven }}-${{ hashFiles('compose.yaml', '**/pom.xml', '**/*.java') }} + key: maven-${{ matrix.jdk }}-${{ matrix.maven }}-${{ hashFiles('compose.yaml', '**/pom.xml') }} restore-keys: maven-${{ matrix.jdk }}-${{ matrix.maven }}- - name: Execute Docker Build env: @@ -88,22 +88,23 @@ jobs: matrix: include: - arch: AMD64 - jdk: 11 - macos: 13 + jdk: 17 + macos: 15-intel - arch: AArch64 - jdk: 11 + jdk: 17 macos: latest steps: - - name: Set up Java - uses: actions/setup-java@v4 - with: - distribution: 'temurin' - java-version: ${{ matrix.jdk }} - name: Checkout Arrow - uses: actions/checkout@v4 + uses: actions/checkout@v7 with: fetch-depth: 0 submodules: recursive + - name: Set up Java + uses: actions/setup-java@v5 + with: + distribution: 'temurin' + java-version: ${{ matrix.jdk }} + cache: 'maven' - name: Build shell: bash env: @@ -123,18 +124,19 @@ jobs: strategy: fail-fast: false matrix: - jdk: [11] + jdk: [17] steps: - - name: Set up Java - uses: actions/setup-java@v4 - with: - java-version: ${{ matrix.jdk }} - distribution: 'temurin' - name: Checkout Arrow - uses: actions/checkout@v4 + uses: actions/checkout@v7 with: fetch-depth: 0 submodules: recursive + - name: Set up Java + uses: actions/setup-java@v5 + with: + java-version: ${{ matrix.jdk }} + distribution: 'temurin' + cache: 'maven' - name: Build shell: bash env: @@ -145,60 +147,3 @@ jobs: env: DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: ci/scripts/test.sh . build jni - - integration: - name: AMD64 integration - runs-on: ubuntu-latest - timeout-minutes: 60 - steps: - - name: Checkout Arrow - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - fetch-depth: 0 - repository: apache/arrow - submodules: recursive - - name: Checkout Arrow Rust - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - repository: apache/arrow-rs - path: rust - - name: Checkout Arrow nanoarrow - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - repository: apache/arrow-nanoarrow - path: nanoarrow - - name: Checkout Arrow Go - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - repository: apache/arrow-go - path: go - - name: Checkout Arrow Java - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - path: java - - name: Free up disk space - run: | - ci/scripts/util_free_space.sh - - name: Cache Docker Volumes - uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3 - with: - path: .docker - key: integration-conda-${{ hashFiles('cpp/**') }} - restore-keys: integration-conda- - - name: Setup Python - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 - with: - python-version: 3.12 - - name: Setup Archery - run: pip install -e dev/archery[docker] - - name: Execute Docker Build - run: | - source ci/scripts/util_enable_core_dumps.sh - archery docker run \ - -e ARCHERY_DEFAULT_BRANCH=main \ - -e ARCHERY_INTEGRATION_TARGET_IMPLEMENTATIONS=java \ - -e ARCHERY_INTEGRATION_WITH_GO=1 \ - -e ARCHERY_INTEGRATION_WITH_JAVA=1 \ - -e ARCHERY_INTEGRATION_WITH_NANOARROW=1 \ - -e ARCHERY_INTEGRATION_WITH_RUST=1 \ - conda-integration diff --git a/.gitignore b/.gitignore index b57597af47..17d1d43ae1 100644 --- a/.gitignore +++ b/.gitignore @@ -7,10 +7,12 @@ .buildpath .checkstyle .classpath +.cursor/ .factorypath .idea/ .project .settings/ +.vscode/ /*-build/ /.mvn/.develocity/ /apache-arrow-java-* diff --git a/.mvn/extensions.xml b/.mvn/extensions.xml index 1dd32d62d6..74482cb2c4 100644 --- a/.mvn/extensions.xml +++ b/.mvn/extensions.xml @@ -23,11 +23,11 @@ com.gradle develocity-maven-extension - 1.23.1 + 2.5.0 com.gradle common-custom-user-data-maven-extension - 2.0.1 + 2.3.0 diff --git a/Brewfile b/Brewfile index af6bd65615..2c47a38af5 100644 --- a/Brewfile +++ b/Brewfile @@ -15,5 +15,5 @@ # specific language governing permissions and limitations # under the License. -brew "openjdk@11" +brew "openjdk@17" brew "sccache" diff --git a/README.md b/README.md index c46b61c49e..0196536514 100644 --- a/README.md +++ b/README.md @@ -23,11 +23,11 @@ The following guides explain the fundamental data structures used in the Java implementation of Apache Arrow. -- [ValueVector](https://arrow.apache.org/docs/java/vector.html) is an abstraction that is used to store a sequence of values having the same type in an individual column. -- [VectorSchemaRoot](https://arrow.apache.org/docs/java/vector_schema_root.html) is a container that can hold multiple vectors based on a schema. -- The [Reading/Writing IPC formats](https://arrow.apache.org/docs/java/ipc.html) guide explains how to stream record batches as well as serializing record batches to files. +- [ValueVector](https://arrow.apache.org/java/current/vector.html) is an abstraction that is used to store a sequence of values having the same type in an individual column. +- [VectorSchemaRoot](https://arrow.apache.org/java/current/vector_schema_root.html#vectorschemaroot) is a container that can hold multiple vectors based on a schema. +- The [Reading/Writing IPC formats](https://arrow.apache.org/java/current/ipc.html) guide explains how to stream record batches as well as serializing record batches to files. -Generated javadoc documentation is available [here](https://arrow.apache.org/docs/java/). +Generated javadoc documentation is available [here](https://arrow.apache.org/java/current/). ## Building from source @@ -50,7 +50,7 @@ a version of your choosing. $ flatc --version flatc version 25.1.24 -$ grep "dep.fbs.version" java/pom.xml +$ grep "dep.fbs.version" pom.xml 25.1.24 ``` @@ -60,10 +60,10 @@ $ grep "dep.fbs.version" java/pom.xml cd $ARROW_HOME # remove the existing files -rm -rf java/format/src +rm -rf format/src # regenerate from the .fbs files -flatc --java -o java/format/src/main/java format/*.fbs +flatc --java -o format/src/main/java arrow-format/*.fbs # prepend license header mvn spotless:apply -pl :arrow-format @@ -93,7 +93,7 @@ conflicting or duplicate fields set this JVM flag or use the correct static cons ## Java Code Style Guide -Arrow Java follows the Google style guide [here][3] with the following +Arrow Java follows the [Google Java Style Guide](http://google.github.io/styleguide/javaguide.html) with the following differences: * Imports are grouped, from top to bottom, in this order: static imports, @@ -119,12 +119,12 @@ following command run in the project root directory: mvn -Dlogback.configurationFile=file: ``` -See [Logback Configuration][1] for more details. +See [Logback Configuration](https://logback.qos.ch/manual/configuration.html) for more details. ## Integration Tests Integration tests which require more time or more memory can be run by activating -the `integration-tests` profile. This activates the [maven failsafe][4] plugin +the `integration-tests` profile. This activates the [Maven Failsafe](https://maven.apache.org/surefire/maven-failsafe-plugin/) plugin and any class prefixed with `IT` will be run during the testing phase. The integration tests currently require a larger amount of memory (>4GB) and time to complete. To activate the profile: @@ -132,8 +132,3 @@ the profile: ```bash mvn -Pintegration-tests ``` - -[1]: https://logback.qos.ch/manual/configuration.html -[2]: https://github.com/apache/arrow/blob/main/cpp/README.md -[3]: http://google.github.io/styleguide/javaguide.html -[4]: https://maven.apache.org/surefire/maven-failsafe-plugin/ diff --git a/adapter/avro/pom.xml b/adapter/avro/pom.xml index 827d19f2a2..4f7f90d7a9 100644 --- a/adapter/avro/pom.xml +++ b/adapter/avro/pom.xml @@ -23,7 +23,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT ../../pom.xml diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java index 87b594af9e..e09b99f670 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java @@ -17,10 +17,14 @@ package org.apache.arrow.adapter.avro; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; +import java.util.regex.Pattern; import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; import org.apache.arrow.adapter.avro.producers.AvroBooleanProducer; import org.apache.arrow.adapter.avro.producers.AvroBytesProducer; +import org.apache.arrow.adapter.avro.producers.AvroEnumProducer; import org.apache.arrow.adapter.avro.producers.AvroFixedSizeBinaryProducer; import org.apache.arrow.adapter.avro.producers.AvroFixedSizeListProducer; import org.apache.arrow.adapter.avro.producers.AvroFloat2Producer; @@ -41,6 +45,7 @@ import org.apache.arrow.adapter.avro.producers.AvroUint8Producer; import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer; +import org.apache.arrow.adapter.avro.producers.DictionaryDecodingProducer; import org.apache.arrow.adapter.avro.producers.Producer; import org.apache.arrow.adapter.avro.producers.logical.AvroDateDayProducer; import org.apache.arrow.adapter.avro.producers.logical.AvroDateMilliProducer; @@ -59,6 +64,7 @@ import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecProducer; import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampSecTzProducer; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; @@ -96,11 +102,14 @@ import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.Text; import org.apache.avro.LogicalType; import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; @@ -162,17 +171,29 @@ public class ArrowToAvroUtils { * may be nullable. Record types must contain at least one child field and cannot contain multiple * fields with the same name * + *

String fields that are dictionary-encoded will be represented as an Avro enum, so long as + * all the values meet the restrictions on Avro enums (non-null, valid identifiers). Other data + * types that are dictionary encoded, or string fields that do not meet the avro requirements, + * will be output as their decoded type. + * * @param arrowFields The arrow fields used to generate the Avro schema * @param typeName Name of the top level Avro record type * @param namespace Namespace of the top level Avro record type + * @param dictionaries A dictionary provider is required if any fields use dictionary encoding * @return An Avro record schema for the given list of fields, with the specified name and * namespace */ public static Schema createAvroSchema( - List arrowFields, String typeName, String namespace) { + List arrowFields, String typeName, String namespace, DictionaryProvider dictionaries) { SchemaBuilder.RecordBuilder assembler = SchemaBuilder.record(typeName).namespace(namespace); - return buildRecordSchema(assembler, arrowFields, namespace); + return buildRecordSchema(assembler, arrowFields, namespace, dictionaries); + } + + /** Overload provided for convenience, sets dictionaries = null. */ + public static Schema createAvroSchema( + List arrowFields, String typeName, String namespace) { + return createAvroSchema(arrowFields, typeName, namespace, null); } /** Overload provided for convenience, sets namespace = null. */ @@ -185,61 +206,83 @@ public static Schema createAvroSchema(List arrowFields) { return createAvroSchema(arrowFields, GENERIC_RECORD_TYPE_NAME); } + /** + * Overload provided for convenience, sets name = GENERIC_RECORD_TYPE_NAME and namespace = null. + */ + public static Schema createAvroSchema(List arrowFields, DictionaryProvider dictionaries) { + return createAvroSchema(arrowFields, GENERIC_RECORD_TYPE_NAME, null, dictionaries); + } + private static T buildRecordSchema( - SchemaBuilder.RecordBuilder builder, List fields, String namespace) { + SchemaBuilder.RecordBuilder builder, + List fields, + String namespace, + DictionaryProvider dictionaries) { if (fields.isEmpty()) { throw new IllegalArgumentException("Record field must have at least one child field"); } SchemaBuilder.FieldAssembler assembler = builder.namespace(namespace).fields(); for (Field field : fields) { - assembler = buildFieldSchema(assembler, field, namespace); + assembler = buildFieldSchema(assembler, field, namespace, dictionaries); } return assembler.endRecord(); } private static SchemaBuilder.FieldAssembler buildFieldSchema( - SchemaBuilder.FieldAssembler assembler, Field field, String namespace) { + SchemaBuilder.FieldAssembler assembler, + Field field, + String namespace, + DictionaryProvider dictionaries) { return assembler .name(field.getName()) - .type(buildTypeSchema(SchemaBuilder.builder(), field, namespace)) + .type(buildTypeSchema(SchemaBuilder.builder(), field, namespace, dictionaries)) .noDefault(); } private static T buildTypeSchema( - SchemaBuilder.TypeBuilder builder, Field field, String namespace) { + SchemaBuilder.TypeBuilder builder, + Field field, + String namespace, + DictionaryProvider dictionaries) { // Nullable unions need special handling, since union types cannot be directly nested if (field.getType().getTypeID() == ArrowType.ArrowTypeID.Union) { boolean unionNullable = field.getChildren().stream().anyMatch(Field::isNullable); if (unionNullable) { SchemaBuilder.UnionAccumulator union = builder.unionOf().nullType(); - return addTypesToUnion(union, field.getChildren(), namespace); + return addTypesToUnion(union, field.getChildren(), namespace, dictionaries); } else { Field headType = field.getChildren().get(0); List tailTypes = field.getChildren().subList(1, field.getChildren().size()); SchemaBuilder.UnionAccumulator union = - buildBaseTypeSchema(builder.unionOf(), headType, namespace); - return addTypesToUnion(union, tailTypes, namespace); + buildBaseTypeSchema(builder.unionOf(), headType, namespace, dictionaries); + return addTypesToUnion(union, tailTypes, namespace, dictionaries); } } else if (field.isNullable()) { - return buildBaseTypeSchema(builder.nullable(), field, namespace); + return buildBaseTypeSchema(builder.nullable(), field, namespace, dictionaries); } else { - return buildBaseTypeSchema(builder, field, namespace); + return buildBaseTypeSchema(builder, field, namespace, dictionaries); } } private static T buildArraySchema( - SchemaBuilder.ArrayBuilder builder, Field listField, String namespace) { + SchemaBuilder.ArrayBuilder builder, + Field listField, + String namespace, + DictionaryProvider dictionaries) { if (listField.getChildren().size() != 1) { throw new IllegalArgumentException("List field must have exactly one child field"); } Field itemField = listField.getChildren().get(0); - return buildTypeSchema(builder.items(), itemField, namespace); + return buildTypeSchema(builder.items(), itemField, namespace, dictionaries); } private static T buildMapSchema( - SchemaBuilder.MapBuilder builder, Field mapField, String namespace) { + SchemaBuilder.MapBuilder builder, + Field mapField, + String namespace, + DictionaryProvider dictionaries) { if (mapField.getChildren().size() != 1) { throw new IllegalArgumentException("Map field must have exactly one child field"); } @@ -253,11 +296,14 @@ private static T buildMapSchema( throw new IllegalArgumentException( "Map keys must be of type string and cannot be nullable for conversion to Avro"); } - return buildTypeSchema(builder.values(), valueField, namespace); + return buildTypeSchema(builder.values(), valueField, namespace, dictionaries); } private static T buildBaseTypeSchema( - SchemaBuilder.BaseTypeBuilder builder, Field field, String namespace) { + SchemaBuilder.BaseTypeBuilder builder, + Field field, + String namespace, + DictionaryProvider dictionaries) { ArrowType.ArrowTypeID typeID = field.getType().getTypeID(); @@ -269,6 +315,33 @@ private static T buildBaseTypeSchema( return builder.booleanType(); case Int: + if (field.getDictionary() != null) { + if (dictionaries == null) { + throw new IllegalArgumentException( + "Field references a dictionary but no dictionaries were provided: " + + field.getName()); + } + Dictionary dictionary = dictionaries.lookup(field.getDictionary().getId()); + if (dictionary == null) { + throw new IllegalArgumentException( + "Field references a dictionary that does not exist: " + + field.getName() + + ", dictionary ID = " + + field.getDictionary().getId()); + } + if (dictionaryIsValidEnum(dictionary)) { + String[] symbols = dictionarySymbols(dictionary); + return builder.enumeration(field.getName()).symbols(symbols); + } else { + Field decodedField = + new Field( + field.getName(), + dictionary.getVector().getField().getFieldType(), + dictionary.getVector().getField().getChildren()); + return buildBaseTypeSchema(builder, decodedField, namespace, dictionaries); + } + } + ArrowType.Int intType = (ArrowType.Int) field.getType(); if (intType.getBitWidth() > 32 || (intType.getBitWidth() == 32 && !intType.getIsSigned())) { return builder.longType(); @@ -328,7 +401,7 @@ private static T buildBaseTypeSchema( String childNamespace = namespace == null ? field.getName() : namespace + "." + field.getName(); return buildRecordSchema( - builder.record(field.getName()), field.getChildren(), childNamespace); + builder.record(field.getName()), field.getChildren(), childNamespace, dictionaries); case List: case FixedSizeList: @@ -339,13 +412,13 @@ private static T buildBaseTypeSchema( new Field("item", itemField.getFieldType(), itemField.getChildren()); Field safeListField = new Field(field.getName(), field.getFieldType(), List.of(safeItemField)); - return buildArraySchema(builder.array(), safeListField, namespace); + return buildArraySchema(builder.array(), safeListField, namespace, dictionaries); } else { - return buildArraySchema(builder.array(), field, namespace); + return buildArraySchema(builder.array(), field, namespace, dictionaries); } case Map: - return buildMapSchema(builder.map(), field, namespace); + return buildMapSchema(builder.map(), field, namespace, dictionaries); default: throw new IllegalArgumentException( @@ -354,9 +427,12 @@ private static T buildBaseTypeSchema( } private static T addTypesToUnion( - SchemaBuilder.UnionAccumulator accumulator, List unionFields, String namespace) { + SchemaBuilder.UnionAccumulator accumulator, + List unionFields, + String namespace, + DictionaryProvider dictionaries) { for (var field : unionFields) { - accumulator = buildBaseTypeSchema(accumulator.and(), field, namespace); + accumulator = buildBaseTypeSchema(accumulator.and(), field, namespace, dictionaries); } return accumulator.endUnion(); } @@ -373,30 +449,88 @@ private static LogicalType timestampLogicalType(ArrowType.Timestamp timestampTyp } } + private static boolean dictionaryIsValidEnum(Dictionary dictionary) { + + if (dictionary.getVectorType().getTypeID() != ArrowType.ArrowTypeID.Utf8) { + return false; + } + + VarCharVector vector = (VarCharVector) dictionary.getVector(); + Set symbols = new HashSet<>(); + + for (int i = 0; i < vector.getValueCount(); i++) { + if (vector.isNull(i)) { + return false; + } + Text text = vector.getObject(i); + if (text == null) { + return false; + } + String symbol = text.toString(); + if (!ENUM_REGEX.matcher(symbol).matches()) { + return false; + } + if (symbols.contains(symbol)) { + return false; + } + symbols.add(symbol); + } + + return true; + } + + private static String[] dictionarySymbols(Dictionary dictionary) { + + VarCharVector vector = (VarCharVector) dictionary.getVector(); + String[] symbols = new String[vector.getValueCount()]; + + for (int i = 0; i < vector.getValueCount(); i++) { + Text text = vector.getObject(i); + // This should never happen if dictionaryIsValidEnum() succeeded + if (text == null) { + throw new IllegalArgumentException("Illegal null value in enum"); + } + symbols[i] = text.toString(); + } + + return symbols; + } + + private static final Pattern ENUM_REGEX = Pattern.compile("^[A-Za-z_][A-Za-z0-9_]*$"); + /** * Create a composite Avro producer for a set of field vectors (typically the root set of a VSR). * * @param vectors The vectors that will be used to produce Avro data * @return The resulting composite Avro producer */ - public static CompositeAvroProducer createCompositeProducer(List vectors) { + public static CompositeAvroProducer createCompositeProducer( + List vectors, DictionaryProvider dictionaries) { List> producers = new ArrayList<>(vectors.size()); for (FieldVector vector : vectors) { - BaseAvroProducer producer = createProducer(vector); + BaseAvroProducer producer = createProducer(vector, dictionaries); producers.add(producer); } return new CompositeAvroProducer(producers); } - private static BaseAvroProducer createProducer(FieldVector vector) { + /** Overload provided for convenience, sets dictionaries = null. */ + public static CompositeAvroProducer createCompositeProducer(List vectors) { + + return createCompositeProducer(vectors, null); + } + + private static BaseAvroProducer createProducer( + FieldVector vector, DictionaryProvider dictionaries) { boolean nullable = vector.getField().isNullable(); - return createProducer(vector, nullable); + return createProducer(vector, nullable, dictionaries); } - private static BaseAvroProducer createProducer(FieldVector vector, boolean nullable) { + private static BaseAvroProducer createProducer( + FieldVector vector, boolean nullable, DictionaryProvider dictionaries) { Preconditions.checkNotNull(vector, "Arrow vector object can't be null"); @@ -405,10 +539,34 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu // Avro understands nullable types as a union of type | null // Most nullable fields in a VSR will not be unions, so provide a special wrapper if (nullable && minorType != Types.MinorType.UNION) { - final BaseAvroProducer innerProducer = createProducer(vector, false); + final BaseAvroProducer innerProducer = createProducer(vector, false, dictionaries); return new AvroNullableProducer<>(innerProducer); } + if (vector.getField().getDictionary() != null) { + if (dictionaries == null) { + throw new IllegalArgumentException( + "Field references a dictionary but no dictionaries were provided: " + + vector.getField().getName()); + } + Dictionary dictionary = dictionaries.lookup(vector.getField().getDictionary().getId()); + if (dictionary == null) { + throw new IllegalArgumentException( + "Field references a dictionary that does not exist: " + + vector.getField().getName() + + ", dictionary ID = " + + vector.getField().getDictionary().getId()); + } + // If a field is dictionary-encoded but cannot be represented as an Avro enum, + // then decode it before writing + if (dictionaryIsValidEnum(dictionary)) { + return new AvroEnumProducer((BaseIntVector) vector); + } else { + BaseAvroProducer dictProducer = createProducer(dictionary.getVector(), false, null); + return new DictionaryDecodingProducer<>((BaseIntVector) vector, dictProducer); + } + } + switch (minorType) { case NULL: return new AvroNullProducer((NullVector) vector); @@ -486,21 +644,23 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu Producer[] childProducers = new Producer[childVectors.size()]; for (int i = 0; i < childVectors.size(); i++) { FieldVector childVector = childVectors.get(i); - childProducers[i] = createProducer(childVector, childVector.getField().isNullable()); + childProducers[i] = + createProducer(childVector, childVector.getField().isNullable(), dictionaries); } return new AvroStructProducer(structVector, childProducers); case LIST: ListVector listVector = (ListVector) vector; FieldVector itemVector = listVector.getDataVector(); - Producer itemProducer = createProducer(itemVector, itemVector.getField().isNullable()); + Producer itemProducer = + createProducer(itemVector, itemVector.getField().isNullable(), dictionaries); return new AvroListProducer(listVector, itemProducer); case FIXED_SIZE_LIST: FixedSizeListVector fixedListVector = (FixedSizeListVector) vector; FieldVector fixedItemVector = fixedListVector.getDataVector(); Producer fixedItemProducer = - createProducer(fixedItemVector, fixedItemVector.getField().isNullable()); + createProducer(fixedItemVector, fixedItemVector.getField().isNullable(), dictionaries); return new AvroFixedSizeListProducer(fixedListVector, fixedItemProducer); case MAP: @@ -514,7 +674,7 @@ private static BaseAvroProducer createProducer(FieldVector vector, boolean nu FieldVector valueVector = entryVector.getChildrenFromFields().get(1); Producer keyProducer = new AvroStringProducer(keyVector); Producer valueProducer = - createProducer(valueVector, valueVector.getField().isNullable()); + createProducer(valueVector, valueVector.getField().isNullable(), dictionaries); Producer entryProducer = new AvroStructProducer(entryVector, new Producer[] {keyProducer, valueProducer}); return new AvroMapProducer(mapVector, entryProducer); diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java index aedef7732e..a6e77e4050 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java @@ -1071,8 +1071,8 @@ private static FieldType createFieldType( } private static String convertAliases(Set aliases) { - JsonStringArrayList jsonList = new JsonStringArrayList(); - aliases.stream().forEach(a -> jsonList.add(a)); + JsonStringArrayList jsonList = new JsonStringArrayList(aliases.size()); + jsonList.addAll(aliases); return jsonList.toString(); } } diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java index 4123370061..e82fdc36fb 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowVectorIterator.java @@ -17,13 +17,14 @@ package org.apache.arrow.adapter.avro; import java.io.EOFException; -import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.stream.Collectors; import org.apache.arrow.adapter.avro.consumers.CompositeAvroConsumer; +import org.apache.arrow.adapter.avro.consumers.Consumer; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.util.ValueVectorUtility; @@ -75,9 +76,11 @@ public static AvroToArrowVectorIterator create( private void initialize() { // create consumers compositeConsumer = AvroToArrowUtils.createCompositeConsumer(schema, config); - List vectors = new ArrayList<>(); - compositeConsumer.getConsumers().forEach(c -> vectors.add(c.getVector())); - List fields = vectors.stream().map(t -> t.getField()).collect(Collectors.toList()); + List vectors = + compositeConsumer.getConsumers().stream() + .map(Consumer::getVector) + .collect(Collectors.toList()); + List fields = vectors.stream().map(ValueVector::getField).collect(Collectors.toList()); VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 0); rootSchema = root.getSchema(); diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java index 068566493e..eebfb7d241 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java @@ -17,22 +17,22 @@ package org.apache.arrow.adapter.avro.producers; import java.io.IOException; -import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.BaseIntVector; import org.apache.avro.io.Encoder; /** - * Producer that produces enum values from a dictionary-encoded {@link IntVector}, writes data to an - * Avro encoder. + * Producer that produces enum values from a dictionary-encoded {@link BaseIntVector}, writes data + * to an Avro encoder. */ -public class AvroEnumProducer extends BaseAvroProducer { +public class AvroEnumProducer extends BaseAvroProducer { /** Instantiate an AvroEnumProducer. */ - public AvroEnumProducer(IntVector vector) { + public AvroEnumProducer(BaseIntVector vector) { super(vector); } @Override public void produce(Encoder encoder) throws IOException { - encoder.writeEnum(vector.get(currentIndex++)); + encoder.writeEnum((int) vector.getValueAsLong(currentIndex++)); } } diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java new file mode 100644 index 0000000000..afeba08511 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/DictionaryDecodingProducer.java @@ -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.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.BaseIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that decodes values from a dictionary-encoded {@link FieldVector}, writes the resulting + * values to an Avro encoder. + * + * @param Type of the underlying dictionary vector + */ +public class DictionaryDecodingProducer + extends BaseAvroProducer { + + private final Producer dictProducer; + + /** Instantiate a DictionaryDecodingProducer. */ + public DictionaryDecodingProducer(BaseIntVector indexVector, Producer dictProducer) { + super(indexVector); + this.dictProducer = dictProducer; + } + + @Override + public void produce(Encoder encoder) throws IOException { + int dicIndex = (int) vector.getValueAsLong(currentIndex++); + dictProducer.setPosition(dicIndex); + dictProducer.produce(encoder); + } +} diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java index 2d70b45021..6d66ee9d45 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java @@ -76,10 +76,14 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryEncoder; +import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.JsonStringArrayList; @@ -2817,4 +2821,81 @@ record = datumReader.read(record, decoder); } } } + + @Test + public void testWriteDictEnumEncoded() throws Exception { + + BufferAllocator allocator = new RootAllocator(); + + // Create a dictionary + FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector dictionaryVector = + new VarCharVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, "apple".getBytes()); + dictionaryVector.set(1, "banana".getBytes()); + dictionaryVector.set(2, "cherry".getBytes()); + dictionaryVector.setValueCount(3); + + Dictionary dictionary = + new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + // Field definition + FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector stringVector = + new VarCharVector(new Field("enumField", stringField, null), allocator); + stringVector.allocateNew(10); + stringVector.setSafe(0, "apple".getBytes()); + stringVector.setSafe(1, "banana".getBytes()); + stringVector.setSafe(2, "cherry".getBytes()); + stringVector.setSafe(3, "cherry".getBytes()); + stringVector.setSafe(4, "apple".getBytes()); + stringVector.setSafe(5, "banana".getBytes()); + stringVector.setSafe(6, "apple".getBytes()); + stringVector.setSafe(7, "cherry".getBytes()); + stringVector.setSafe(8, "banana".getBytes()); + stringVector.setSafe(9, "apple".getBytes()); + stringVector.setValueCount(10); + + IntVector encodedVector = (IntVector) DictionaryEncoder.encode(stringVector, dictionary); + + // Set up VSR + List vectors = Arrays.asList(encodedVector); + int rowCount = 10; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + File dataFile = new File(TMP, "testWriteEnumEncoded.avro"); + + // Write an AVRO block using the producer classes + try (FileOutputStream fos = new FileOutputStream(dataFile)) { + BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null); + CompositeAvroProducer producer = + ArrowToAvroUtils.createCompositeProducer(vectors, dictionaries); + for (int row = 0; row < rowCount; row++) { + producer.produce(encoder); + } + encoder.flush(); + } + + // Set up reading the AVRO block as a GenericRecord + Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields(), dictionaries); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + GenericRecord record = null; + + // Read and check values + for (int row = 0; row < rowCount; row++) { + record = datumReader.read(record, decoder); + // Values read from Avro should be the decoded enum values + assertEquals(stringVector.getObject(row).toString(), record.get("enumField").toString()); + } + } + } + } } diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java index d3e12e763a..d5e0357a8c 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java @@ -20,11 +20,18 @@ import java.util.Arrays; import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.UnionMode; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.avro.LogicalTypes; @@ -1389,4 +1396,126 @@ public void testConvertUnionTypes() { Schema.Type.STRING, schema.getField("nullableDenseUnionField").schema().getTypes().get(3).getType()); } + + @Test + public void testWriteDictEnumEncoded() { + + BufferAllocator allocator = new RootAllocator(); + + // Create a dictionary + FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector dictionaryVector = + new VarCharVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, "apple".getBytes()); + dictionaryVector.set(1, "banana".getBytes()); + dictionaryVector.set(2, "cherry".getBytes()); + dictionaryVector.setValueCount(3); + + Dictionary dictionary = + new Dictionary( + dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + List fields = + Arrays.asList( + new Field( + "enumField", + new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(1, schema.getFields().size()); + + Schema.Field enumField = schema.getField("enumField"); + + assertEquals(Schema.Type.ENUM, enumField.schema().getType()); + assertEquals(3, enumField.schema().getEnumSymbols().size()); + assertEquals("apple", enumField.schema().getEnumSymbols().get(0)); + assertEquals("banana", enumField.schema().getEnumSymbols().get(1)); + assertEquals("cherry", enumField.schema().getEnumSymbols().get(2)); + } + + @Test + public void testWriteDictEnumInvalid() { + + BufferAllocator allocator = new RootAllocator(); + + // Create a dictionary + FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector dictionaryVector = + new VarCharVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, "passion fruit".getBytes()); + dictionaryVector.set(1, "banana".getBytes()); + dictionaryVector.set(2, "cherry".getBytes()); + dictionaryVector.setValueCount(3); + + Dictionary dictionary = + new Dictionary( + dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + List fields = + Arrays.asList( + new Field( + "enumField", + new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null), + null)); + + // Dictionary field contains values that are not valid enums + // Should be decoded and output as a string field + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(1, schema.getFields().size()); + + Schema.Field enumField = schema.getField("enumField"); + assertEquals(Schema.Type.STRING, enumField.schema().getType()); + } + + @Test + public void testWriteDictEnumInvalid2() { + + BufferAllocator allocator = new RootAllocator(); + + // Create a dictionary + FieldType dictionaryField = new FieldType(false, new ArrowType.Int(64, true), null); + BigIntVector dictionaryVector = + new BigIntVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, 123L); + dictionaryVector.set(1, 456L); + dictionaryVector.set(2, 789L); + dictionaryVector.setValueCount(3); + + Dictionary dictionary = + new Dictionary( + dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + List fields = + Arrays.asList( + new Field( + "enumField", + new FieldType(false, new ArrowType.Int(8, true), dictionary.getEncoding(), null), + null)); + + // Dictionary field encodes LONG values rather than STRING + // Should be doecded and output as a LONG field + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(1, schema.getFields().size()); + + Schema.Field enumField = schema.getField("enumField"); + assertEquals(Schema.Type.LONG, enumField.schema().getType()); + } } diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java index 85e6a960b0..ceaf59aa72 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java @@ -52,6 +52,7 @@ import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.TimeStampNanoTZVector; import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; @@ -60,10 +61,14 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryEncoder; +import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.avro.Schema; @@ -78,16 +83,21 @@ public class RoundTripDataTest { @TempDir public static File TMP; - private static AvroToArrowConfig basicConfig(BufferAllocator allocator) { - return new AvroToArrowConfig(allocator, 1000, null, Collections.emptySet(), false); + private static AvroToArrowConfig basicConfig( + BufferAllocator allocator, DictionaryProvider.MapDictionaryProvider dictionaries) { + return new AvroToArrowConfig(allocator, 1000, dictionaries, Collections.emptySet(), false); } private static VectorSchemaRoot readDataFile( - Schema schema, File dataFile, BufferAllocator allocator) throws Exception { + Schema schema, + File dataFile, + BufferAllocator allocator, + DictionaryProvider.MapDictionaryProvider dictionaries) + throws Exception { try (FileInputStream fis = new FileInputStream(dataFile)) { BinaryDecoder decoder = new DecoderFactory().directBinaryDecoder(fis, null); - return AvroToArrow.avroToArrow(schema, decoder, basicConfig(allocator)); + return AvroToArrow.avroToArrow(schema, decoder, basicConfig(allocator, dictionaries)); } } @@ -95,11 +105,22 @@ private static void roundTripTest( VectorSchemaRoot root, BufferAllocator allocator, File dataFile, int rowCount) throws Exception { + roundTripTest(root, allocator, dataFile, rowCount, null); + } + + private static void roundTripTest( + VectorSchemaRoot root, + BufferAllocator allocator, + File dataFile, + int rowCount, + DictionaryProvider dictionaries) + throws Exception { + // Write an AVRO block using the producer classes try (FileOutputStream fos = new FileOutputStream(dataFile)) { BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null); CompositeAvroProducer producer = - ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors()); + ArrowToAvroUtils.createCompositeProducer(root.getFieldVectors(), dictionaries); for (int row = 0; row < rowCount; row++) { producer.produce(encoder); } @@ -107,10 +128,14 @@ private static void roundTripTest( } // Generate AVRO schema - Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields()); + Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields(), dictionaries); + + DictionaryProvider.MapDictionaryProvider roundTripDictionaries = + new DictionaryProvider.MapDictionaryProvider(); // Read back in and compare - try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator)) { + try (VectorSchemaRoot roundTrip = + readDataFile(schema, dataFile, allocator, roundTripDictionaries)) { assertEquals(root.getSchema(), roundTrip.getSchema()); assertEquals(rowCount, roundTrip.getRowCount()); @@ -119,6 +144,21 @@ private static void roundTripTest( for (int row = 0; row < rowCount; row++) { assertEquals(root.getVector(0).getObject(row), roundTrip.getVector(0).getObject(row)); } + + if (dictionaries != null) { + for (long id : dictionaries.getDictionaryIds()) { + Dictionary originalDictionary = dictionaries.lookup(id); + Dictionary roundTripDictionary = roundTripDictionaries.lookup(id); + assertEquals( + originalDictionary.getVector().getValueCount(), + roundTripDictionary.getVector().getValueCount()); + for (int j = 0; j < originalDictionary.getVector().getValueCount(); j++) { + assertEquals( + originalDictionary.getVector().getObject(j), + roundTripDictionary.getVector().getObject(j)); + } + } + } } } @@ -141,7 +181,7 @@ private static void roundTripByteArrayTest( Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields()); // Read back in and compare - try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator)) { + try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator, null)) { assertEquals(root.getSchema(), roundTrip.getSchema()); assertEquals(rowCount, roundTrip.getRowCount()); @@ -1603,4 +1643,58 @@ public void testRoundTripNullableStructs() throws Exception { roundTripTest(root, allocator, dataFile, rowCount); } } + + @Test + public void testRoundTripEnum() throws Exception { + + BufferAllocator allocator = new RootAllocator(); + + // Create a dictionary + FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector dictionaryVector = + new VarCharVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, "apple".getBytes()); + dictionaryVector.set(1, "banana".getBytes()); + dictionaryVector.set(2, "cherry".getBytes()); + dictionaryVector.setValueCount(3); + + // For simplicity, ensure the index type matches what will be decoded during Avro enum decoding + Dictionary dictionary = + new Dictionary( + dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + // Field definition + FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector stringVector = + new VarCharVector(new Field("enumField", stringField, null), allocator); + stringVector.allocateNew(10); + stringVector.setSafe(0, "apple".getBytes()); + stringVector.setSafe(1, "banana".getBytes()); + stringVector.setSafe(2, "cherry".getBytes()); + stringVector.setSafe(3, "cherry".getBytes()); + stringVector.setSafe(4, "apple".getBytes()); + stringVector.setSafe(5, "banana".getBytes()); + stringVector.setSafe(6, "apple".getBytes()); + stringVector.setSafe(7, "cherry".getBytes()); + stringVector.setSafe(8, "banana".getBytes()); + stringVector.setSafe(9, "apple".getBytes()); + stringVector.setValueCount(10); + + TinyIntVector encodedVector = + (TinyIntVector) DictionaryEncoder.encode(stringVector, dictionary); + + // Set up VSR + List vectors = Arrays.asList(encodedVector); + int rowCount = 10; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + File dataFile = new File(TMP, "testRoundTripEnums.avro"); + + roundTripTest(root, allocator, dataFile, rowCount, dictionaries); + } + } } diff --git a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java index 864e2c8b59..37c0b4d9fe 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java @@ -21,27 +21,50 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.DictionaryEncoding; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.avro.Schema; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class RoundTripSchemaTest { private void doRoundTripTest(List fields) { + doRoundTripTest(fields, null); + } - AvroToArrowConfig config = new AvroToArrowConfig(null, 1, null, Collections.emptySet(), false); + private void doRoundTripTest(List fields, DictionaryProvider dictionaries) { - Schema avroSchema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + DictionaryProvider.MapDictionaryProvider decodeDictionaries = + new DictionaryProvider.MapDictionaryProvider(); + AvroToArrowConfig decodeConfig = + new AvroToArrowConfig(null, 1, decodeDictionaries, Collections.emptySet(), false); + + Schema avroSchema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord", null, dictionaries); org.apache.arrow.vector.types.pojo.Schema arrowSchema = - AvroToArrowUtils.createArrowSchema(avroSchema, config); + AvroToArrowUtils.createArrowSchema(avroSchema, decodeConfig); // Compare string representations - equality not defined for logical types assertEquals(fields, arrowSchema.getFields()); + + for (int i = 0; i < fields.size(); i++) { + Field field = fields.get(i); + Field rtField = arrowSchema.getFields().get(i); + if (field.getDictionary() != null) { + // Dictionary content is not decoded until the data is consumed + Assertions.assertNotNull(rtField.getDictionary()); + } + } } // Schema round trip for primitive types, nullable and non-nullable @@ -440,4 +463,38 @@ public void testRoundTripStructType() { doRoundTripTest(fields); } + + @Test + public void testRoundTripEnumType() { + + BufferAllocator allocator = new RootAllocator(); + + FieldType dictionaryField = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector dictionaryVector = + new VarCharVector(new Field("dictionary", dictionaryField, null), allocator); + + dictionaryVector.allocateNew(3); + dictionaryVector.set(0, "apple".getBytes()); + dictionaryVector.set(1, "banana".getBytes()); + dictionaryVector.set(2, "cherry".getBytes()); + dictionaryVector.setValueCount(3); + + // For simplicity, ensure the index type matches what will be decoded during Avro enum decoding + Dictionary dictionary = + new Dictionary( + dictionaryVector, new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))); + DictionaryProvider dictionaries = new DictionaryProvider.MapDictionaryProvider(dictionary); + + List fields = + Arrays.asList( + new Field( + "enumField", + new FieldType( + true, + new ArrowType.Int(8, true), + new DictionaryEncoding(0L, false, new ArrowType.Int(8, true))), + null)); + + doRoundTripTest(fields, dictionaries); + } } diff --git a/adapter/jdbc/pom.xml b/adapter/jdbc/pom.xml index 2f621d7a05..a8ac19721d 100644 --- a/adapter/jdbc/pom.xml +++ b/adapter/jdbc/pom.xml @@ -23,7 +23,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT ../../pom.xml @@ -59,7 +59,7 @@ under the License. com.h2database h2 - 2.3.232 + 2.4.240 test diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java index 2366116fd0..b8389ee27c 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java @@ -24,7 +24,6 @@ import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.ArrowType; /** Composite consumer which hold all consumers. It manages the consume and cleanup process. */ public class CompositeJdbcConsumer implements JdbcConsumer { @@ -46,9 +45,9 @@ public void consume(ResultSet rs) throws SQLException, IOException { BaseConsumer consumer = (BaseConsumer) consumers[i]; JdbcFieldInfo fieldInfo = new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet); - ArrowType arrowType = consumer.vector.getMinorType().getType(); + throw new JdbcConsumerException( - "Exception while consuming JDBC value", e, fieldInfo, arrowType); + "Exception while consuming JDBC value", e, fieldInfo, consumer.vector.getField()); } else { throw e; } diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java index 04e26d640c..98927f416c 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/exceptions/JdbcConsumerException.java @@ -17,7 +17,7 @@ package org.apache.arrow.adapter.jdbc.consumer.exceptions; import org.apache.arrow.adapter.jdbc.JdbcFieldInfo; -import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; /** * Exception while consuming JDBC data. This exception stores the JdbcFieldInfo for the column and @@ -25,7 +25,7 @@ */ public class JdbcConsumerException extends RuntimeException { final JdbcFieldInfo fieldInfo; - final ArrowType arrowType; + final Field field; /** * Construct JdbcConsumerException with all fields. @@ -33,17 +33,17 @@ public class JdbcConsumerException extends RuntimeException { * @param message error message * @param cause original exception * @param fieldInfo JdbcFieldInfo for the column - * @param arrowType ArrowType for the corresponding vector + * @param field ArrowType for the corresponding vector */ public JdbcConsumerException( - String message, Throwable cause, JdbcFieldInfo fieldInfo, ArrowType arrowType) { + String message, Throwable cause, JdbcFieldInfo fieldInfo, Field field) { super(message, cause); this.fieldInfo = fieldInfo; - this.arrowType = arrowType; + this.field = field; } - public ArrowType getArrowType() { - return this.arrowType; + public Field getField() { + return this.field; } public JdbcFieldInfo getFieldInfo() { diff --git a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java index c7dc9b2791..e5039ccf59 100644 --- a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/ResultSetUtilityTest.java @@ -43,15 +43,19 @@ public void testZeroRowResultSet() throws Exception { .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) .build(); - ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config); - assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use"); - VectorSchemaRoot root = iter.next(); - assertNotNull(root, "VectorSchemaRoot from first next() result should never be null"); - assertEquals( - 0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows"); - assertFalse( - iter.hasNext(), - "hasNext() should return false on empty ResultSets after initial next() call"); + try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { + assertTrue(iter.hasNext(), "Iterator on zero row ResultSet should haveNext() before use"); + VectorSchemaRoot root = iter.next(); + assertNotNull(root, "VectorSchemaRoot from first next() result should never be null"); + assertEquals( + 0, root.getRowCount(), "VectorSchemaRoot from empty ResultSet should have zero rows"); + assertFalse( + iter.hasNext(), + "hasNext() should return false on empty ResultSets after initial next() call"); + if (!reuseVectorSchemaRoot) { + root.close(); + } + } } } } diff --git a/adapter/orc/pom.xml b/adapter/orc/pom.xml index e3ae7d5163..50a9b3a603 100644 --- a/adapter/orc/pom.xml +++ b/adapter/orc/pom.xml @@ -23,7 +23,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT ../../pom.xml @@ -61,7 +61,7 @@ under the License. org.apache.orc orc-core - 2.1.1 + 2.3.0 test diff --git a/algorithm/pom.xml b/algorithm/pom.xml index 898c2605b6..24adcefa6f 100644 --- a/algorithm/pom.xml +++ b/algorithm/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-algorithm Arrow Algorithms diff --git a/arrow-format/FlightSql.proto b/arrow-format/FlightSql.proto index 3568d851cb..b1dc57b33b 100644 --- a/arrow-format/FlightSql.proto +++ b/arrow-format/FlightSql.proto @@ -1212,6 +1212,7 @@ message CommandGetDbSchemas { * - ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE - "1" indicates if the column is case-sensitive, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_READ_ONLY - "1" indicates if the column is read only, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_SEARCHABLE - "1" indicates if the column is searchable via WHERE clause, "0" otherwise. + * - ARROW:FLIGHT:SQL:REMARKS - A comment describing the column. * The returned data should be ordered by catalog_name, db_schema_name, table_name, then table_type, followed by table_schema if requested. */ message CommandGetTables { @@ -1549,6 +1550,11 @@ message ActionCreatePreparedStatementResult { // If the query provided contained parameters, parameter_schema contains the // schema of the expected parameters. It should be an IPC-encapsulated Schema, as described in Schema.fbs. bytes parameter_schema = 3; + + // When set to true, the query should be executed with CommandPreparedStatementUpdate, + // when set to false, the query should be executed with CommandPreparedStatementQuery. + // If not set, the client can choose how to execute the query. + optional bool is_update = 4; } /* @@ -1678,6 +1684,7 @@ message ActionEndSavepointRequest { * - ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE - "1" indicates if the column is case-sensitive, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_READ_ONLY - "1" indicates if the column is read only, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_SEARCHABLE - "1" indicates if the column is searchable via WHERE clause, "0" otherwise. + * - ARROW:FLIGHT:SQL:REMARKS - A comment describing the column. * - GetFlightInfo: execute the query. */ message CommandStatementQuery { @@ -1703,6 +1710,7 @@ message CommandStatementQuery { * - ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE - "1" indicates if the column is case-sensitive, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_READ_ONLY - "1" indicates if the column is read only, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_SEARCHABLE - "1" indicates if the column is searchable via WHERE clause, "0" otherwise. + * - ARROW:FLIGHT:SQL:REMARKS - A comment describing the column. * - GetFlightInfo: execute the query. * - DoPut: execute the query. */ @@ -1739,6 +1747,7 @@ message TicketStatementQuery { * - ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE - "1" indicates if the column is case-sensitive, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_READ_ONLY - "1" indicates if the column is read only, "0" otherwise. * - ARROW:FLIGHT:SQL:IS_SEARCHABLE - "1" indicates if the column is searchable via WHERE clause, "0" otherwise. + * - ARROW:FLIGHT:SQL:REMARKS - A comment describing the column. * * If the schema is retrieved after parameter values have been bound with DoPut, then the server should account * for the parameters when determining the schema. diff --git a/arrow-variant/pom.xml b/arrow-variant/pom.xml new file mode 100644 index 0000000000..e578626dd4 --- /dev/null +++ b/arrow-variant/pom.xml @@ -0,0 +1,51 @@ + + + + 4.0.0 + + org.apache.arrow + arrow-java-root + 20.0.0-SNAPSHOT + + arrow-variant + Arrow Variant + Arrow Variant type support. + + + + org.apache.arrow + arrow-memory-core + + + org.apache.arrow + arrow-vector + + + org.apache.parquet + parquet-variant + ${dep.parquet.version} + + + org.apache.arrow + arrow-memory-unsafe + test + + + diff --git a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java b/arrow-variant/src/main/java/module-info.java similarity index 69% rename from vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java rename to arrow-variant/src/main/java/module-info.java index 207b0951a7..da94173969 100644 --- a/vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java +++ b/arrow-variant/src/main/java/module-info.java @@ -14,10 +14,15 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.arrow.vector.holder; -import org.apache.arrow.vector.holders.ExtensionHolder; +@SuppressWarnings("requires-automatic") +module org.apache.arrow.variant { + exports org.apache.arrow.variant; + exports org.apache.arrow.variant.extension; + exports org.apache.arrow.variant.impl; + exports org.apache.arrow.variant.holders; -public class UuidHolder extends ExtensionHolder { - public byte[] value; + requires org.apache.arrow.memory.core; + requires org.apache.arrow.vector; + requires parquet.variant; } diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java b/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java new file mode 100644 index 0000000000..fa05cdd93f --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/Variant.java @@ -0,0 +1,217 @@ +/* + * 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.arrow.variant; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.Objects; +import java.util.UUID; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.variant.holders.NullableVariantHolder; + +/** + * Wrapper around parquet-variant's Variant implementation. + * + *

This wrapper exists to isolate the parquet-variant dependency from Arrow's public API, + * allowing the vector module to expose variant functionality without requiring users to depend on + * parquet-variant directly. It also ensures that nested variant values (from arrays and objects) + * are consistently wrapped. + */ +public class Variant { + + private final org.apache.parquet.variant.Variant delegate; + + /** Creates a Variant from raw metadata and value byte arrays. */ + public Variant(byte[] metadata, byte[] value) { + this.delegate = new org.apache.parquet.variant.Variant(value, metadata); + } + + /** Creates a Variant by copying data from ArrowBuf instances. */ + public Variant( + ArrowBuf metadataBuffer, + int metadataStart, + int metadataEnd, + ArrowBuf valueBuffer, + int valueStart, + int valueEnd) { + byte[] metadata = new byte[metadataEnd - metadataStart]; + byte[] value = new byte[valueEnd - valueStart]; + metadataBuffer.getBytes(metadataStart, metadata); + valueBuffer.getBytes(valueStart, value); + this.delegate = new org.apache.parquet.variant.Variant(value, metadata); + } + + private Variant(org.apache.parquet.variant.Variant delegate) { + this.delegate = delegate; + } + + /** Constructs a Variant from a NullableVariantHolder. */ + public Variant(NullableVariantHolder holder) { + this( + holder.metadataBuffer, + holder.metadataStart, + holder.metadataEnd, + holder.valueBuffer, + holder.valueStart, + holder.valueEnd); + } + + public ByteBuffer getValueBuffer() { + return delegate.getValueBuffer(); + } + + public ByteBuffer getMetadataBuffer() { + return delegate.getMetadataBuffer(); + } + + public boolean getBoolean() { + return delegate.getBoolean(); + } + + public byte getByte() { + return delegate.getByte(); + } + + public short getShort() { + return delegate.getShort(); + } + + public int getInt() { + return delegate.getInt(); + } + + public long getLong() { + return delegate.getLong(); + } + + public double getDouble() { + return delegate.getDouble(); + } + + public BigDecimal getDecimal() { + return delegate.getDecimal(); + } + + public float getFloat() { + return delegate.getFloat(); + } + + public ByteBuffer getBinary() { + return delegate.getBinary(); + } + + public UUID getUUID() { + return delegate.getUUID(); + } + + public String getString() { + return delegate.getString(); + } + + public Type getType() { + return Type.fromParquet(delegate.getType()); + } + + public int numObjectElements() { + return delegate.numObjectElements(); + } + + public Variant getFieldByKey(String key) { + org.apache.parquet.variant.Variant result = delegate.getFieldByKey(key); + return result != null ? wrap(result) : null; + } + + public ObjectField getFieldAtIndex(int idx) { + org.apache.parquet.variant.Variant.ObjectField field = delegate.getFieldAtIndex(idx); + return new ObjectField(field.key, wrap(field.value)); + } + + public int numArrayElements() { + return delegate.numArrayElements(); + } + + public Variant getElementAtIndex(int index) { + org.apache.parquet.variant.Variant result = delegate.getElementAtIndex(index); + return result != null ? wrap(result) : null; + } + + private static Variant wrap(org.apache.parquet.variant.Variant parquetVariant) { + return new Variant(parquetVariant); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + Variant variant = (Variant) o; + return delegate.getMetadataBuffer().equals(variant.delegate.getMetadataBuffer()) + && delegate.getValueBuffer().equals(variant.delegate.getValueBuffer()); + } + + @Override + public int hashCode() { + return Objects.hash(delegate.getMetadataBuffer(), delegate.getValueBuffer()); + } + + @Override + public String toString() { + return "Variant{type=" + getType() + '}'; + } + + public enum Type { + OBJECT, + ARRAY, + NULL, + BOOLEAN, + BYTE, + SHORT, + INT, + LONG, + STRING, + DOUBLE, + DECIMAL4, + DECIMAL8, + DECIMAL16, + DATE, + TIMESTAMP_TZ, + TIMESTAMP_NTZ, + FLOAT, + BINARY, + TIME, + TIMESTAMP_NANOS_TZ, + TIMESTAMP_NANOS_NTZ, + UUID; + + static Type fromParquet(org.apache.parquet.variant.Variant.Type parquetType) { + return Type.valueOf(parquetType.name()); + } + } + + public static final class ObjectField { + public final String key; + public final Variant value; + + public ObjectField(String key, Variant value) { + this.key = key; + this.value = value; + } + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java new file mode 100644 index 0000000000..3deb70cdc0 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantType.java @@ -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.arrow.variant.extension; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.variant.impl.VariantWriterImpl; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.FieldType; + +/** + * Arrow extension type for Parquet + * Variant binary encoding. The type itself does not support shredded variant data. + */ +public final class VariantType extends ExtensionType { + + public static final VariantType INSTANCE = new VariantType(); + + public static final String EXTENSION_NAME = "parquet.variant"; + + static { + ExtensionTypeRegistry.register(INSTANCE); + } + + private VariantType() {} + + @Override + public ArrowType storageType() { + return ArrowType.Struct.INSTANCE; + } + + @Override + public String extensionName() { + return EXTENSION_NAME; + } + + @Override + public boolean extensionEquals(ExtensionType other) { + return other instanceof VariantType; + } + + @Override + public String serialize() { + return ""; + } + + @Override + public ArrowType deserialize(ArrowType storageType, String serializedData) { + if (!storageType.equals(this.storageType())) { + throw new UnsupportedOperationException( + "Cannot construct VariantType from underlying type " + storageType); + } + return INSTANCE; + } + + @Override + public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { + return new VariantVector(name, allocator); + } + + @Override + public boolean isComplex() { + // The type itself is not complex meaning we need separate functions to convert/extract + // different types. + // Meanwhile, the containing vector is complex in terms of containing multiple values (metadata + // and value) + return false; + } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new VariantWriterImpl((VariantVector) vector); + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java new file mode 100644 index 0000000000..1bbf1a6bdb --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/extension/VariantVector.java @@ -0,0 +1,348 @@ +/* + * 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.arrow.variant.extension; + +import java.nio.ByteBuffer; +import java.util.List; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.variant.holders.VariantHolder; +import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.complex.AbstractStructVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.Binary; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.CallBack; +import org.apache.arrow.vector.util.TransferPair; + +/** + * Arrow vector for storing {@link VariantType} values. + * + *

Stores semi-structured data (like JSON) as metadata + value binary pairs, allowing + * type-flexible columnar storage within Arrow's type system. + */ +public class VariantVector extends ExtensionTypeVector { + + public static final String METADATA_VECTOR_NAME = "metadata"; + public static final String VALUE_VECTOR_NAME = "value"; + + private final Field rootField; + + /** + * Constructs a new VariantVector with the given name and allocator. + * + * @param name the name of the vector + * @param allocator the buffer allocator for memory management + */ + public VariantVector(String name, BufferAllocator allocator) { + super( + name, + allocator, + new StructVector( + name, + allocator, + FieldType.nullable(ArrowType.Struct.INSTANCE), + null, + AbstractStructVector.ConflictPolicy.CONFLICT_ERROR, + false)); + rootField = createVariantField(name); + ((FieldVector) this.getUnderlyingVector()) + .initializeChildrenFromFields(rootField.getChildren()); + } + + /** + * Creates a new VariantVector with the given name. The Variant Field schema has to be the same + * everywhere, otherwise ArrowBuffer loading might fail during serialization/deserialization and + * schema mismatches can occur. This includes CompleteType's VARIANT and VARIANT_REQUIRED types. + */ + public static Field createVariantField(String name) { + return new Field( + name, new FieldType(true, VariantType.INSTANCE, null), createVariantChildFields()); + } + + /** + * Creates the child fields for the VariantVector. Metadata vector will be index 0 and value + * vector will be index 1. + */ + public static List createVariantChildFields() { + return List.of( + new Field(METADATA_VECTOR_NAME, new FieldType(false, Binary.INSTANCE, null), null), + new Field(VALUE_VECTOR_NAME, new FieldType(false, Binary.INSTANCE, null), null)); + } + + @Override + public void initializeChildrenFromFields(List children) { + // No-op, as children are initialized in the constructor + } + + @Override + public Field getField() { + return rootField; + } + + public VarBinaryVector getMetadataVector() { + return getUnderlyingVector().getChild(METADATA_VECTOR_NAME, VarBinaryVector.class); + } + + public VarBinaryVector getValueVector() { + return getUnderlyingVector().getChild(VALUE_VECTOR_NAME, VarBinaryVector.class); + } + + @Override + public TransferPair makeTransferPair(ValueVector target) { + return new VariantTransferPair(this, (VariantVector) target); + } + + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator) { + return new VariantTransferPair(this, new VariantVector(field.getName(), allocator)); + } + + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(field, allocator); + } + + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator) { + return new VariantTransferPair(this, new VariantVector(ref, allocator)); + } + + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(ref, allocator); + } + + @Override + public TransferPair getTransferPair(BufferAllocator allocator) { + return getTransferPair(this.getField().getName(), allocator); + } + + @Override + public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { + getUnderlyingVector() + .copyFrom(fromIndex, thisIndex, ((VariantVector) from).getUnderlyingVector()); + } + + @Override + public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { + getUnderlyingVector() + .copyFromSafe(fromIndex, thisIndex, ((VariantVector) from).getUnderlyingVector()); + } + + @Override + public Object getObject(int index) { + if (isNull(index)) { + return null; + } + VarBinaryVector metadataVector = getMetadataVector(); + VarBinaryVector valueVector = getValueVector(); + + int metadataStart = metadataVector.getStartOffset(index); + int metadataEnd = metadataVector.getEndOffset(index); + int valueStart = valueVector.getStartOffset(index); + int valueEnd = valueVector.getEndOffset(index); + + return new Variant( + metadataVector.getDataBuffer(), + metadataStart, + metadataEnd, + valueVector.getDataBuffer(), + valueStart, + valueEnd); + } + + /** + * Retrieves the variant value at the specified index into the provided holder. + * + * @param index the index of the value to retrieve + * @param holder the holder to populate with the variant data + */ + public void get(int index, NullableVariantHolder holder) { + if (isNull(index)) { + holder.isSet = 0; + } else { + holder.isSet = 1; + VarBinaryVector metadataVector = getMetadataVector(); + VarBinaryVector valueVector = getValueVector(); + assert !metadataVector.isNull(index) && !valueVector.isNull(index); + + holder.metadataStart = metadataVector.getStartOffset(index); + holder.metadataEnd = metadataVector.getEndOffset(index); + holder.metadataBuffer = metadataVector.getDataBuffer(); + holder.valueStart = valueVector.getStartOffset(index); + holder.valueEnd = valueVector.getEndOffset(index); + holder.valueBuffer = valueVector.getDataBuffer(); + } + } + + /** + * Retrieves the variant value at the specified index into the provided non-nullable holder. + * + * @param index the index of the value to retrieve + * @param holder the holder to populate with the variant data + */ + public void get(int index, VariantHolder holder) { + VarBinaryVector metadataVector = getMetadataVector(); + VarBinaryVector valueVector = getValueVector(); + assert !metadataVector.isNull(index) && !valueVector.isNull(index); + + holder.metadataStart = metadataVector.getStartOffset(index); + holder.metadataEnd = metadataVector.getEndOffset(index); + holder.metadataBuffer = metadataVector.getDataBuffer(); + holder.valueStart = valueVector.getStartOffset(index); + holder.valueEnd = valueVector.getEndOffset(index); + holder.valueBuffer = valueVector.getDataBuffer(); + } + + /** + * Sets the variant value at the specified index from the provided holder. + * + * @param index the index at which to set the value + * @param holder the holder containing the variant data to set + */ + public void set(int index, VariantHolder holder) { + BitVectorHelper.setBit(getUnderlyingVector().getValidityBuffer(), index); + getMetadataVector() + .set(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer); + getValueVector().set(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer); + } + + /** + * Sets the variant value at the specified index from the provided nullable holder. + * + * @param index the index at which to set the value + * @param holder the nullable holder containing the variant data to set + */ + public void set(int index, NullableVariantHolder holder) { + BitVectorHelper.setValidityBit(getUnderlyingVector().getValidityBuffer(), index, holder.isSet); + if (holder.isSet == 0) { + return; + } + getMetadataVector() + .set(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer); + getValueVector().set(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer); + } + + /** + * Sets the variant value at the specified index from the provided holder, with bounds checking. + * + * @param index the index at which to set the value + * @param holder the holder containing the variant data to set + */ + public void setSafe(int index, VariantHolder holder) { + getUnderlyingVector().setIndexDefined(index); + getMetadataVector() + .setSafe(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer); + getValueVector().setSafe(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer); + } + + /** + * Sets the variant value at the specified index from the provided nullable holder, with bounds + * checking. + * + * @param index the index at which to set the value + * @param holder the nullable holder containing the variant data to set + */ + public void setSafe(int index, NullableVariantHolder holder) { + if (holder.isSet == 0) { + getUnderlyingVector().setNull(index); + return; + } + getUnderlyingVector().setIndexDefined(index); + getMetadataVector() + .setSafe(index, 1, holder.metadataStart, holder.metadataEnd, holder.metadataBuffer); + getValueVector().setSafe(index, 1, holder.valueStart, holder.valueEnd, holder.valueBuffer); + } + + /** Sets the value at the given index from the provided Variant. */ + public void setSafe(int index, Variant variant) { + ByteBuffer metadataBuffer = variant.getMetadataBuffer(); + ByteBuffer valueBuffer = variant.getValueBuffer(); + int metadataLength = metadataBuffer.remaining(); + int valueLength = valueBuffer.remaining(); + try (ArrowBuf metaBuf = getAllocator().buffer(metadataLength); + ArrowBuf valBuf = getAllocator().buffer(valueLength)) { + metaBuf.setBytes(0, metadataBuffer.duplicate()); + valBuf.setBytes(0, valueBuffer.duplicate()); + getUnderlyingVector().setIndexDefined(index); + getMetadataVector().setSafe(index, 1, 0, metadataLength, metaBuf); + getValueVector().setSafe(index, 1, 0, valueLength, valBuf); + } + } + + @Override + protected FieldReader getReaderImpl() { + return new org.apache.arrow.variant.impl.VariantReaderImpl(this); + } + + @Override + public int hashCode(int index) { + return hashCode(index, null); + } + + @Override + public int hashCode(int index, ArrowBufHasher hasher) { + return getUnderlyingVector().hashCode(index, hasher); + } + + /** + * VariantTransferPair is a transfer pair for VariantVector. It transfers the metadata and value + * together using the underlyingVector's transfer pair. + */ + protected static class VariantTransferPair implements TransferPair { + private final TransferPair pair; + private final VariantVector from; + private final VariantVector to; + + public VariantTransferPair(VariantVector from, VariantVector to) { + this.from = from; + this.to = to; + this.pair = from.getUnderlyingVector().makeTransferPair((to).getUnderlyingVector()); + } + + @Override + public void transfer() { + pair.transfer(); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + pair.splitAndTransfer(startIndex, length); + } + + @Override + public ValueVector getTo() { + return to; + } + + @Override + public void copyValueSafe(int from, int to) { + pair.copyValueSafe(from, to); + } + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java new file mode 100644 index 0000000000..b78d4a2013 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/NullableVariantHolder.java @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.variant.holders; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.variant.extension.VariantType; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; + +@SuppressWarnings("checkstyle:VisibilityModifier") +public final class NullableVariantHolder extends ExtensionHolder { + + public int isSet; + public int metadataStart; + public int metadataEnd; + public ArrowBuf metadataBuffer; + public int valueStart; + public int valueEnd; + public ArrowBuf valueBuffer; + + public NullableVariantHolder() {} + + @Override + public boolean equals(Object obj) { + throw new UnsupportedOperationException(); + } + + @Override + public int hashCode() { + throw new UnsupportedOperationException(); + } + + @Override + public String toString() { + throw new UnsupportedOperationException(); + } + + @Override + public ArrowType type() { + return VariantType.INSTANCE; + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java similarity index 51% rename from vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java rename to arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java index 68029b1df5..e3947ac439 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java @@ -14,34 +14,43 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.arrow.vector.complex.impl; +package org.apache.arrow.variant.holders; -import java.nio.ByteBuffer; -import java.util.UUID; -import org.apache.arrow.vector.UuidVector; -import org.apache.arrow.vector.holder.UuidHolder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.variant.extension.VariantType; import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; -public class UuidWriterImpl extends AbstractExtensionTypeWriter { +@SuppressWarnings("checkstyle:VisibilityModifier") +public final class VariantHolder extends ExtensionHolder { - public UuidWriterImpl(UuidVector vector) { - super(vector); + public final int isSet = 1; + public int metadataStart; + public int metadataEnd; + public ArrowBuf metadataBuffer; + public int valueStart; + public int valueEnd; + public ArrowBuf valueBuffer; + + public VariantHolder() {} + + @Override + public boolean equals(Object obj) { + throw new UnsupportedOperationException(); + } + + @Override + public int hashCode() { + throw new UnsupportedOperationException(); } @Override - public void writeExtension(Object value) { - UUID uuid = (UUID) value; - ByteBuffer bb = ByteBuffer.allocate(16); - bb.putLong(uuid.getMostSignificantBits()); - bb.putLong(uuid.getLeastSignificantBits()); - vector.setSafe(getPosition(), bb.array()); - vector.setValueCount(getPosition() + 1); + public String toString() { + throw new UnsupportedOperationException(); } @Override - public void write(ExtensionHolder holder) { - UuidHolder uuidHolder = (UuidHolder) holder; - vector.setSafe(getPosition(), uuidHolder.value); - vector.setValueCount(getPosition() + 1); + public ArrowType type() { + return VariantType.INSTANCE; } } diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java new file mode 100644 index 0000000000..1645529c0c --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/NullableVariantHolderReaderImpl.java @@ -0,0 +1,69 @@ +/* + * 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.arrow.variant.impl; + +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.vector.complex.impl.AbstractFieldReader; +import org.apache.arrow.vector.types.Types; + +public class NullableVariantHolderReaderImpl extends AbstractFieldReader { + private final NullableVariantHolder holder; + + public NullableVariantHolderReaderImpl(NullableVariantHolder holder) { + this.holder = holder; + } + + @Override + public int size() { + throw new UnsupportedOperationException("You can't call size on a Holder value reader."); + } + + @Override + public boolean next() { + throw new UnsupportedOperationException("You can't call next on a single value reader."); + } + + @Override + public void setPosition(int index) { + throw new UnsupportedOperationException("You can't call setPosition on a single value reader."); + } + + @Override + public Types.MinorType getMinorType() { + return Types.MinorType.EXTENSIONTYPE; + } + + @Override + public boolean isSet() { + return holder.isSet == 1; + } + + /** + * Reads the variant holder data into the provided holder. + * + * @param h the holder to read into + */ + public void read(NullableVariantHolder h) { + h.metadataStart = this.holder.metadataStart; + h.metadataEnd = this.holder.metadataEnd; + h.metadataBuffer = this.holder.metadataBuffer; + h.valueStart = this.holder.valueStart; + h.valueEnd = this.holder.valueEnd; + h.valueBuffer = this.holder.valueBuffer; + h.isSet = this.isSet() ? 1 : 0; + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java new file mode 100644 index 0000000000..670104b7d1 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantReaderImpl.java @@ -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.arrow.variant.impl; + +import org.apache.arrow.variant.extension.VariantVector; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.variant.holders.VariantHolder; +import org.apache.arrow.vector.complex.impl.AbstractFieldReader; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.pojo.Field; + +public class VariantReaderImpl extends AbstractFieldReader { + private final VariantVector vector; + + public VariantReaderImpl(VariantVector vector) { + this.vector = vector; + } + + @Override + public Types.MinorType getMinorType() { + return this.vector.getMinorType(); + } + + @Override + public Field getField() { + return this.vector.getField(); + } + + @Override + public boolean isSet() { + return !this.vector.isNull(this.idx()); + } + + @Override + public void read(ExtensionHolder holder) { + if (holder instanceof VariantHolder) { + vector.get(idx(), (VariantHolder) holder); + } else if (holder instanceof NullableVariantHolder) { + vector.get(idx(), (NullableVariantHolder) holder); + } else { + throw new IllegalArgumentException( + "Unsupported holder type for VariantReader: " + holder.getClass()); + } + } + + public void read(VariantHolder h) { + this.vector.get(this.idx(), h); + } + + public void read(NullableVariantHolder h) { + this.vector.get(this.idx(), h); + } + + @Override + public Object readObject() { + return this.vector.getObject(this.idx()); + } +} diff --git a/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java new file mode 100644 index 0000000000..266ddb75d2 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/impl/VariantWriterImpl.java @@ -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.arrow.variant.impl; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.extension.VariantVector; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.variant.holders.VariantHolder; +import org.apache.arrow.vector.complex.impl.AbstractExtensionTypeWriter; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** + * Writer implementation for VARIANT extension type vectors. + * + *

This writer handles writing variant data to a {@link VariantVector}. It accepts both {@link + * VariantHolder} and {@link NullableVariantHolder} objects containing metadata and value buffers + * and writes them to the appropriate position in the vector. + */ +public class VariantWriterImpl extends AbstractExtensionTypeWriter { + + private static final String UNSUPPORTED_TYPE_TEMPLATE = "Unsupported type for Variant: %s"; + + /** + * Constructs a new VariantWriterImpl for the given vector. + * + * @param vector the variant vector to write to + */ + public VariantWriterImpl(VariantVector vector) { + super(vector); + } + + /** + * Writes an extension type or variant value to the vector. + * + *

This method handles {@link ExtensionHolder} by delegating to {@link #write(ExtensionHolder)} + * and {@link Variant} by delegating to {@link #writeVariant(Variant)}. + * + * @param object the object to write, must be an {@link ExtensionHolder} or {@link Variant} + * @throws IllegalArgumentException if the object is not an {@link ExtensionHolder} or {@link + * Variant} + */ + @Override + public void writeExtension(Object object) { + if (object instanceof ExtensionHolder) { + write((ExtensionHolder) object); + } else if (object instanceof Variant) { + writeVariant((Variant) object); + } else { + throw new IllegalArgumentException( + String.format(UNSUPPORTED_TYPE_TEMPLATE, object.getClass().getName())); + } + } + + private void writeVariant(Variant variant) { + java.nio.ByteBuffer metadataBuffer = variant.getMetadataBuffer(); + java.nio.ByteBuffer valueBuffer = variant.getValueBuffer(); + int metadataLength = metadataBuffer.remaining(); + int valueLength = valueBuffer.remaining(); + try (ArrowBuf metadataBuf = vector.getAllocator().buffer(metadataLength); + ArrowBuf valueBuf = vector.getAllocator().buffer(valueLength)) { + metadataBuf.setBytes(0, metadataBuffer.duplicate()); + valueBuf.setBytes(0, valueBuffer.duplicate()); + NullableVariantHolder holder = new NullableVariantHolder(); + holder.isSet = 1; + holder.metadataBuffer = metadataBuf; + holder.metadataStart = 0; + holder.metadataEnd = metadataLength; + holder.valueBuffer = valueBuf; + holder.valueStart = 0; + holder.valueEnd = valueLength; + vector.setSafe(getPosition(), holder); + vector.setValueCount(getPosition() + 1); + } + } + + @Override + public void writeExtension(Object value, ArrowType type) { + writeExtension(value); + } + + /** + * Writes a variant holder to the vector at the current position. + * + *

The holder can be either a {@link VariantHolder} (non-nullable, always set) or a {@link + * NullableVariantHolder} (nullable, may be null). The data is written using {@link + * VariantVector#setSafe(int, NullableVariantHolder)} which handles buffer allocation and copying. + * + * @param extensionHolder the variant holder to write, must be a {@link VariantHolder} or {@link + * NullableVariantHolder} + * @throws IllegalArgumentException if the holder is neither a {@link VariantHolder} nor a {@link + * NullableVariantHolder} + */ + @Override + public void write(ExtensionHolder extensionHolder) { + if (extensionHolder instanceof VariantHolder) { + vector.setSafe(getPosition(), (VariantHolder) extensionHolder); + } else if (extensionHolder instanceof NullableVariantHolder) { + vector.setSafe(getPosition(), (NullableVariantHolder) extensionHolder); + } else { + throw new IllegalArgumentException( + String.format(UNSUPPORTED_TYPE_TEMPLATE, extensionHolder.getClass().getName())); + } + vector.setValueCount(getPosition() + 1); + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java b/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java new file mode 100644 index 0000000000..bc46a68616 --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/TestVariant.java @@ -0,0 +1,439 @@ +/* + * 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.arrow.variant; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.UUID; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.parquet.variant.VariantBuilder; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestVariant { + + private BufferAllocator allocator; + + @BeforeEach + void beforeEach() { + allocator = new RootAllocator(); + } + + @AfterEach + void afterEach() { + allocator.close(); + } + + static Variant buildVariant(VariantBuilder builder) { + org.apache.parquet.variant.Variant parquetVariant = builder.build(); + ByteBuffer valueBuf = parquetVariant.getValueBuffer(); + ByteBuffer metaBuf = parquetVariant.getMetadataBuffer(); + byte[] valueBytes = new byte[valueBuf.remaining()]; + byte[] metaBytes = new byte[metaBuf.remaining()]; + valueBuf.get(valueBytes); + metaBuf.get(metaBytes); + return new Variant(metaBytes, valueBytes); + } + + public static Variant variantString(String value) { + VariantBuilder builder = new VariantBuilder(); + builder.appendString(value); + return buildVariant(builder); + } + + @Test + void testConstructionWithArrowBuf() { + VariantBuilder builder = new VariantBuilder(); + builder.appendInt(42); + Variant source = buildVariant(builder); + int metaLen = source.getMetadataBuffer().remaining(); + int valueLen = source.getValueBuffer().remaining(); + + try (ArrowBuf metadataArrowBuf = allocator.buffer(metaLen + 2); + ArrowBuf valueArrowBuf = allocator.buffer(valueLen + 3)) { + metadataArrowBuf.setBytes(2, source.getMetadataBuffer()); + valueArrowBuf.setBytes(3, source.getValueBuffer()); + + Variant variant = + new Variant(metadataArrowBuf, 2, 2 + metaLen, valueArrowBuf, 3, 3 + valueLen); + + assertEquals(Variant.Type.INT, variant.getType()); + assertEquals(42, variant.getInt()); + } + } + + @Test + void testNullType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendNull(); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.NULL, variant.getType()); + } + + @Test + void testBooleanType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendBoolean(true); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.BOOLEAN, variant.getType()); + assertTrue(variant.getBoolean()); + + builder = new VariantBuilder(); + builder.appendBoolean(false); + variant = buildVariant(builder); + + assertEquals(Variant.Type.BOOLEAN, variant.getType()); + assertFalse(variant.getBoolean()); + } + + @Test + void testByteType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendByte((byte) 42); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.BYTE, variant.getType()); + assertEquals((byte) 42, variant.getByte()); + } + + @Test + void testShortType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendShort((short) 1234); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.SHORT, variant.getType()); + assertEquals((short) 1234, variant.getShort()); + } + + @Test + void testIntType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendInt(123456); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.INT, variant.getType()); + assertEquals(123456, variant.getInt()); + } + + @Test + void testLongType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendLong(9876543210L); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.LONG, variant.getType()); + assertEquals(9876543210L, variant.getLong()); + } + + @Test + void testFloatType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendFloat(3.14f); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.FLOAT, variant.getType()); + assertEquals(3.14f, variant.getFloat(), 0.001f); + } + + @Test + void testDoubleType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendDouble(3.14159265359); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.DOUBLE, variant.getType()); + assertEquals(3.14159265359, variant.getDouble(), 0.0000001); + } + + @Test + void testStringType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendString("hello world"); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.STRING, variant.getType()); + assertEquals("hello world", variant.getString()); + } + + @Test + void testDecimalType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendDecimal(new BigDecimal("123.456")); + Variant variant = buildVariant(builder); + + assertTrue( + variant.getType() == Variant.Type.DECIMAL4 + || variant.getType() == Variant.Type.DECIMAL8 + || variant.getType() == Variant.Type.DECIMAL16); + assertEquals(new BigDecimal("123.456"), variant.getDecimal()); + } + + @Test + void testBinaryType() { + VariantBuilder builder = new VariantBuilder(); + byte[] data = new byte[] {1, 2, 3, 4, 5}; + builder.appendBinary(ByteBuffer.wrap(data)); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.BINARY, variant.getType()); + ByteBuffer result = variant.getBinary(); + byte[] resultBytes = new byte[result.remaining()]; + result.get(resultBytes); + assertArrayEquals(data, resultBytes); + } + + @Test + void testUuidType() { + VariantBuilder builder = new VariantBuilder(); + UUID uuid = UUID.randomUUID(); + builder.appendUUID(uuid); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.UUID, variant.getType()); + assertEquals(uuid, variant.getUUID()); + } + + @Test + void testDateType() { + VariantBuilder builder = new VariantBuilder(); + int daysSinceEpoch = 19000; + builder.appendDate(daysSinceEpoch); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.DATE, variant.getType()); + } + + @Test + void testTimestampTzType() { + VariantBuilder builder = new VariantBuilder(); + long micros = System.currentTimeMillis() * 1000; + builder.appendTimestampTz(micros); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.TIMESTAMP_TZ, variant.getType()); + } + + @Test + void testTimestampNtzType() { + VariantBuilder builder = new VariantBuilder(); + long micros = System.currentTimeMillis() * 1000; + builder.appendTimestampNtz(micros); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.TIMESTAMP_NTZ, variant.getType()); + } + + @Test + void testTimeType() { + VariantBuilder builder = new VariantBuilder(); + long micros = 12345678L; + builder.appendTime(micros); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.TIME, variant.getType()); + } + + @Test + void testObjectType() { + VariantBuilder builder = new VariantBuilder(); + var objBuilder = builder.startObject(); + objBuilder.appendKey("name"); + objBuilder.appendString("test"); + objBuilder.appendKey("value"); + objBuilder.appendInt(42); + builder.endObject(); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.OBJECT, variant.getType()); + assertEquals(2, variant.numObjectElements()); + + Variant nameField = variant.getFieldByKey("name"); + assertNotNull(nameField); + assertEquals(Variant.Type.STRING, nameField.getType()); + assertEquals("test", nameField.getString()); + + Variant valueField = variant.getFieldByKey("value"); + assertNotNull(valueField); + assertEquals(Variant.Type.INT, valueField.getType()); + assertEquals(42, valueField.getInt()); + + assertNull(variant.getFieldByKey("nonexistent")); + + // Empty object + builder = new VariantBuilder(); + builder.startObject(); + builder.endObject(); + Variant emptyObj = buildVariant(builder); + assertEquals(Variant.Type.OBJECT, emptyObj.getType()); + assertEquals(0, emptyObj.numObjectElements()); + } + + @Test + void testObjectFieldAtIndex() { + VariantBuilder builder = new VariantBuilder(); + var objBuilder = builder.startObject(); + objBuilder.appendKey("alpha"); + objBuilder.appendInt(1); + objBuilder.appendKey("beta"); + objBuilder.appendInt(2); + builder.endObject(); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.OBJECT, variant.getType()); + assertEquals(2, variant.numObjectElements()); + + Variant.ObjectField field0 = variant.getFieldAtIndex(0); + assertNotNull(field0); + assertNotNull(field0.key); + assertNotNull(field0.value); + + Variant.ObjectField field1 = variant.getFieldAtIndex(1); + assertNotNull(field1); + assertNotNull(field1.key); + assertNotNull(field1.value); + } + + @Test + void testArrayType() { + VariantBuilder builder = new VariantBuilder(); + var arrayBuilder = builder.startArray(); + arrayBuilder.appendInt(1); + arrayBuilder.appendInt(2); + arrayBuilder.appendInt(3); + builder.endArray(); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.ARRAY, variant.getType()); + assertEquals(3, variant.numArrayElements()); + + Variant elem0 = variant.getElementAtIndex(0); + assertNotNull(elem0); + assertEquals(Variant.Type.INT, elem0.getType()); + assertEquals(1, elem0.getInt()); + + Variant elem1 = variant.getElementAtIndex(1); + assertEquals(2, elem1.getInt()); + + Variant elem2 = variant.getElementAtIndex(2); + assertEquals(3, elem2.getInt()); + + assertNull(variant.getElementAtIndex(-1)); + assertNull(variant.getElementAtIndex(3)); + + // Empty array + builder = new VariantBuilder(); + builder.startArray(); + builder.endArray(); + Variant emptyArr = buildVariant(builder); + assertEquals(Variant.Type.ARRAY, emptyArr.getType()); + assertEquals(0, emptyArr.numArrayElements()); + } + + @Test + void testNestedStructure() { + VariantBuilder builder = new VariantBuilder(); + var objBuilder = builder.startObject(); + objBuilder.appendKey("items"); + var arrayBuilder = objBuilder.startArray(); + arrayBuilder.appendString("a"); + arrayBuilder.appendString("b"); + objBuilder.endArray(); + builder.endObject(); + Variant variant = buildVariant(builder); + + assertEquals(Variant.Type.OBJECT, variant.getType()); + Variant items = variant.getFieldByKey("items"); + assertNotNull(items); + assertEquals(Variant.Type.ARRAY, items.getType()); + assertEquals(2, items.numArrayElements()); + assertEquals("a", items.getElementAtIndex(0).getString()); + assertEquals("b", items.getElementAtIndex(1).getString()); + } + + @Test + void testEquals() { + VariantBuilder builder1 = new VariantBuilder(); + builder1.appendString("test"); + Variant variant1 = buildVariant(builder1); + + VariantBuilder builder2 = new VariantBuilder(); + builder2.appendString("test"); + Variant variant2 = buildVariant(builder2); + + VariantBuilder builder3 = new VariantBuilder(); + builder3.appendString("different"); + Variant variant3 = buildVariant(builder3); + + assertEquals(variant1, variant1); + assertEquals(variant1, variant2); + assertNotEquals(variant1, variant3); + assertNotEquals(variant1, null); + assertNotEquals(variant1, "not a variant"); + } + + @Test + void testHashCode() { + VariantBuilder builder1 = new VariantBuilder(); + builder1.appendInt(42); + Variant variant1 = buildVariant(builder1); + + VariantBuilder builder2 = new VariantBuilder(); + builder2.appendInt(42); + Variant variant2 = buildVariant(builder2); + + assertEquals(variant1.hashCode(), variant2.hashCode()); + } + + @Test + void testToString() { + VariantBuilder builder = new VariantBuilder(); + builder.appendString("test"); + Variant variant = buildVariant(builder); + + String str = variant.toString(); + assertNotNull(str); + assertTrue(str.contains("type=")); + } + + @Test + void testTypeEnumsMatch() { + for (Variant.Type arrowType : Variant.Type.values()) { + org.apache.parquet.variant.Variant.Type parquetType = + org.apache.parquet.variant.Variant.Type.valueOf(arrowType.name()); + assertEquals(arrowType, Variant.Type.fromParquet(parquetType)); + } + for (org.apache.parquet.variant.Variant.Type parquetType : + org.apache.parquet.variant.Variant.Type.values()) { + Variant.Type arrowType = Variant.Type.valueOf(parquetType.name()); + assertEquals(parquetType.name(), arrowType.name()); + } + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java new file mode 100644 index 0000000000..f3213d523a --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantExtensionType.java @@ -0,0 +1,249 @@ +/* + * 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.arrow.variant.extension; + +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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.File; +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.channels.SeekableByteChannel; +import java.nio.channels.WritableByteChannel; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.nio.file.StandardOpenOption; +import java.util.Collections; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.variant.TestVariant; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.compare.Range; +import org.apache.arrow.vector.compare.RangeEqualsVisitor; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.writer.BaseWriter; +import org.apache.arrow.vector.ipc.ArrowFileReader; +import org.apache.arrow.vector.ipc.ArrowFileWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.VectorBatchAppender; +import org.apache.arrow.vector.validate.ValidateVectorVisitor; +import org.junit.jupiter.api.Test; + +public class TestVariantExtensionType { + + private static void ensureRegistered(ArrowType.ExtensionType type) { + if (ExtensionTypeRegistry.lookup(type.extensionName()) == null) { + ExtensionTypeRegistry.register(type); + } + } + + @Test + public void roundtripVariant() throws IOException { + ensureRegistered(VariantType.INSTANCE); + final Schema schema = + new Schema(Collections.singletonList(Field.nullable("a", VariantType.INSTANCE))); + try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VariantVector vector = (VariantVector) root.getVector("a"); + vector.allocateNew(); + + vector.setSafe(0, TestVariant.variantString("hello")); + vector.setSafe(1, TestVariant.variantString("world")); + vector.setValueCount(2); + root.setRowCount(2); + + final File file = File.createTempFile("varianttest", ".arrow"); + try (final WritableByteChannel channel = + FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE); + final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + + try (final SeekableByteChannel channel = + Files.newByteChannel(Paths.get(file.getAbsolutePath())); + final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { + reader.loadNextBatch(); + final VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); + assertEquals(root.getSchema(), readerRoot.getSchema()); + + final Field field = readerRoot.getSchema().getFields().get(0); + final VariantType expectedType = VariantType.INSTANCE; + assertEquals( + field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), + expectedType.extensionName()); + assertEquals( + field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA), + expectedType.serialize()); + + final ExtensionTypeVector deserialized = + (ExtensionTypeVector) readerRoot.getFieldVectors().get(0); + assertEquals(vector.getValueCount(), deserialized.getValueCount()); + for (int i = 0; i < vector.getValueCount(); i++) { + assertEquals(vector.isNull(i), deserialized.isNull(i)); + if (!vector.isNull(i)) { + assertEquals(vector.getObject(i), deserialized.getObject(i)); + } + } + } + } + } + + @Test + public void readVariantAsUnderlyingType() throws IOException { + ensureRegistered(VariantType.INSTANCE); + final Schema schema = + new Schema(Collections.singletonList(VariantVector.createVariantField("a"))); + try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VariantVector vector = (VariantVector) root.getVector("a"); + vector.allocateNew(); + + vector.setSafe(0, TestVariant.variantString("hello")); + vector.setValueCount(1); + root.setRowCount(1); + + final File file = File.createTempFile("varianttest", ".arrow"); + try (final WritableByteChannel channel = + FileChannel.open(Paths.get(file.getAbsolutePath()), StandardOpenOption.WRITE); + final ArrowFileWriter writer = new ArrowFileWriter(root, null, channel)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + + ExtensionTypeRegistry.unregister(VariantType.INSTANCE); + + try (final SeekableByteChannel channel = + Files.newByteChannel(Paths.get(file.getAbsolutePath())); + final ArrowFileReader reader = new ArrowFileReader(channel, allocator)) { + reader.loadNextBatch(); + VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); + + // Verify schema properties + assertEquals(1, readRoot.getSchema().getFields().size()); + assertEquals("a", readRoot.getSchema().getFields().get(0).getName()); + assertTrue(readRoot.getSchema().getFields().get(0).getType() instanceof ArrowType.Struct); + + // Verify extension metadata is preserved + final Field field = readRoot.getSchema().getFields().get(0); + assertEquals( + VariantType.EXTENSION_NAME, + field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME)); + assertEquals("", field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA)); + + // Verify vector type and row count + assertEquals(1, readRoot.getRowCount()); + FieldVector readVector = readRoot.getVector("a"); + assertEquals(StructVector.class, readVector.getClass()); + + // Verify value count matches + StructVector structVector = (StructVector) readVector; + assertEquals(vector.getValueCount(), structVector.getValueCount()); + + // Verify the underlying data can be accessed from child vectors + VarBinaryVector metadataVector = + structVector.getChild(VariantVector.METADATA_VECTOR_NAME, VarBinaryVector.class); + VarBinaryVector valueVector = + structVector.getChild(VariantVector.VALUE_VECTOR_NAME, VarBinaryVector.class); + assertNotNull(metadataVector); + assertNotNull(valueVector); + assertEquals(1, metadataVector.getValueCount()); + assertEquals(1, valueVector.getValueCount()); + } + } + } + + @Test + public void testVariantVectorCompare() { + VariantType variantType = VariantType.INSTANCE; + ExtensionTypeRegistry.register(variantType); + Variant hello = TestVariant.variantString("hello"); + Variant world = TestVariant.variantString("world"); + try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); + VariantVector a1 = + (VariantVector) + variantType.getNewVector("a", FieldType.nullable(variantType), allocator); + VariantVector a2 = + (VariantVector) + variantType.getNewVector("a", FieldType.nullable(variantType), allocator); + VariantVector bb = + (VariantVector) + variantType.getNewVector("a", FieldType.nullable(variantType), allocator)) { + + ValidateVectorVisitor validateVisitor = new ValidateVectorVisitor(); + validateVisitor.visit(a1, null); + + a1.allocateNew(); + a2.allocateNew(); + bb.allocateNew(); + + a1.setSafe(0, hello); + a1.setSafe(1, world); + a1.setValueCount(2); + + a2.setSafe(0, hello); + a2.setSafe(1, world); + a2.setValueCount(2); + + bb.setSafe(0, world); + bb.setSafe(1, hello); + bb.setValueCount(2); + + Range range = new Range(0, 0, a1.getValueCount()); + RangeEqualsVisitor visitor = new RangeEqualsVisitor(a1, a2); + assertTrue(visitor.rangeEquals(range)); + + visitor = new RangeEqualsVisitor(a1, bb); + assertFalse(visitor.rangeEquals(range)); + + VectorBatchAppender.batchAppend(a1, a2, bb); + assertEquals(6, a1.getValueCount()); + validateVisitor.visit(a1, null); + } + } + + @Test + public void testVariantCopyAsValueThrowsException() { + ensureRegistered(VariantType.INSTANCE); + try (BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); + VariantVector vector = new VariantVector("variant", allocator)) { + vector.allocateNew(); + vector.setSafe(0, TestVariant.variantString("hello")); + vector.setValueCount(1); + + var reader = vector.getReader(); + reader.setPosition(0); + + assertThrows( + IllegalArgumentException.class, () -> reader.copyAsValue((BaseWriter.StructWriter) null)); + } + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java new file mode 100644 index 0000000000..8b6000bc46 --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInListVector.java @@ -0,0 +1,202 @@ +/* + * 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.arrow.variant.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.variant.TestVariant; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListReader; +import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.TransferPair; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestVariantInListVector { + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void terminate() throws Exception { + allocator.close(); + } + + @Test + public void testListVectorWithVariantExtensionType() { + final FieldType type = FieldType.nullable(VariantType.INSTANCE); + try (ListVector inVector = new ListVector("input", allocator, type, null)) { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("bye"); + + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + + writer.setPosition(0); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(VariantType.INSTANCE); + extensionWriter.writeExtension(variant1); + extensionWriter.writeExtension(variant2); + writer.endList(); + inVector.setValueCount(1); + + ArrayList resultSet = (ArrayList) inVector.getObject(0); + assertEquals(2, resultSet.size()); + assertEquals(variant1, resultSet.get(0)); + assertEquals(variant2, resultSet.get(1)); + } + } + + @Test + public void testListVectorReaderForVariantExtensionType() { + try (ListVector inVector = ListVector.empty("input", allocator)) { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("bye"); + + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + + writer.setPosition(0); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(VariantType.INSTANCE); + extensionWriter.writeExtension(variant1); + writer.endList(); + + writer.setPosition(1); + writer.startList(); + extensionWriter.writeExtension(variant2); + extensionWriter.writeExtension(variant2); + writer.endList(); + + inVector.setValueCount(2); + + UnionListReader reader = inVector.getReader(); + reader.setPosition(0); + assertTrue(reader.next()); + FieldReader variantReader = reader.reader(); + NullableVariantHolder resultHolder = new NullableVariantHolder(); + variantReader.read(resultHolder); + assertEquals(variant1, new Variant(resultHolder)); + + reader.setPosition(1); + assertTrue(reader.next()); + variantReader = reader.reader(); + variantReader.read(resultHolder); + assertEquals(variant2, new Variant(resultHolder)); + + assertTrue(reader.next()); + variantReader = reader.reader(); + variantReader.read(resultHolder); + assertEquals(variant2, new Variant(resultHolder)); + } + } + + @Test + public void testCopyFromForVariantExtensionType() { + try (ListVector inVector = ListVector.empty("input", allocator); + ListVector outVector = ListVector.empty("output", allocator)) { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("bye"); + + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + + writer.setPosition(0); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(VariantType.INSTANCE); + extensionWriter.writeExtension(variant1); + writer.endList(); + + writer.setPosition(1); + writer.startList(); + extensionWriter.writeExtension(variant2); + extensionWriter.writeExtension(variant2); + writer.endList(); + + inVector.setValueCount(2); + + outVector.allocateNew(); + outVector.copyFrom(0, 0, inVector); + outVector.copyFrom(1, 1, inVector); + outVector.setValueCount(2); + + ArrayList resultSet0 = (ArrayList) outVector.getObject(0); + assertEquals(1, resultSet0.size()); + assertEquals(variant1, resultSet0.get(0)); + + ArrayList resultSet1 = (ArrayList) outVector.getObject(1); + assertEquals(2, resultSet1.size()); + assertEquals(variant2, resultSet1.get(0)); + assertEquals(variant2, resultSet1.get(1)); + } + } + + @Test + public void testCopyValueSafeForVariantExtensionType() { + try (ListVector inVector = ListVector.empty("input", allocator)) { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("bye"); + + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + + writer.setPosition(0); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(VariantType.INSTANCE); + extensionWriter.writeExtension(variant1); + writer.endList(); + + writer.setPosition(1); + writer.startList(); + extensionWriter.writeExtension(variant2); + extensionWriter.writeExtension(variant2); + writer.endList(); + + inVector.setValueCount(2); + + try (ListVector outVector = (ListVector) inVector.getTransferPair(allocator).getTo()) { + TransferPair tp = inVector.makeTransferPair(outVector); + tp.copyValueSafe(0, 0); + tp.copyValueSafe(1, 1); + outVector.setValueCount(2); + + ArrayList resultSet0 = (ArrayList) outVector.getObject(0); + assertEquals(1, resultSet0.size()); + assertEquals(variant1, resultSet0.get(0)); + + ArrayList resultSet1 = (ArrayList) outVector.getObject(1); + assertEquals(2, resultSet1.size()); + assertEquals(variant2, resultSet1.get(0)); + assertEquals(variant2, resultSet1.get(1)); + } + } + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java new file mode 100644 index 0000000000..dd925810de --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantInMapVector.java @@ -0,0 +1,125 @@ +/* + * 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.arrow.variant.extension; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.variant.TestVariant; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.complex.impl.UnionMapReader; +import org.apache.arrow.vector.complex.impl.UnionMapWriter; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestVariantInMapVector { + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void terminate() { + allocator.close(); + } + + @Test + public void testMapVectorWithVariantExtensionType() { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("world"); + try (final MapVector inVector = MapVector.empty("map", allocator, false)) { + inVector.allocateNew(); + UnionMapWriter writer = inVector.getWriter(); + writer.setPosition(0); + + writer.startMap(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(0); + writer.value().extension(VariantType.INSTANCE).writeExtension(variant1, VariantType.INSTANCE); + writer.endEntry(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(1); + writer.value().extension(VariantType.INSTANCE).writeExtension(variant2, VariantType.INSTANCE); + writer.endEntry(); + writer.endMap(); + + writer.setValueCount(1); + + UnionMapReader mapReader = inVector.getReader(); + mapReader.setPosition(0); + mapReader.next(); + FieldReader variantReader = mapReader.value(); + NullableVariantHolder holder = new NullableVariantHolder(); + variantReader.read(holder); + assertEquals(variant1, new Variant(holder)); + + mapReader.next(); + variantReader = mapReader.value(); + variantReader.read(holder); + assertEquals(variant2, new Variant(holder)); + } + } + + @Test + public void testCopyFromForVariantExtensionType() { + Variant variant1 = TestVariant.variantString("hello"); + Variant variant2 = TestVariant.variantString("world"); + try (final MapVector inVector = MapVector.empty("in", allocator, false); + final MapVector outVector = MapVector.empty("out", allocator, false)) { + inVector.allocateNew(); + UnionMapWriter writer = inVector.getWriter(); + writer.setPosition(0); + + writer.startMap(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(0); + writer.value().extension(VariantType.INSTANCE).writeExtension(variant1, VariantType.INSTANCE); + writer.endEntry(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(1); + writer.value().extension(VariantType.INSTANCE).writeExtension(variant2, VariantType.INSTANCE); + writer.endEntry(); + writer.endMap(); + + writer.setValueCount(1); + outVector.allocateNew(); + outVector.copyFrom(0, 0, inVector); + outVector.setValueCount(1); + + UnionMapReader mapReader = outVector.getReader(); + mapReader.setPosition(0); + mapReader.next(); + FieldReader variantReader = mapReader.value(); + NullableVariantHolder holder = new NullableVariantHolder(); + variantReader.read(holder); + assertEquals(variant1, new Variant(holder)); + + mapReader.next(); + variantReader = mapReader.value(); + variantReader.read(holder); + assertEquals(variant2, new Variant(holder)); + } + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java new file mode 100644 index 0000000000..017e71224b --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantType.java @@ -0,0 +1,308 @@ +/* + * 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.arrow.variant.extension; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +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 java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.dictionary.DictionaryProvider; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestVariantType { + BufferAllocator allocator; + + @BeforeEach + void beforeEach() { + allocator = new RootAllocator(); + } + + @AfterEach + void afterEach() { + allocator.close(); + } + + @Test + void testConstants() { + assertNotNull(VariantType.INSTANCE); + } + + @Test + void testStorageType() { + VariantType type = VariantType.INSTANCE; + assertEquals(ArrowType.Struct.INSTANCE, type.storageType()); + assertInstanceOf(ArrowType.Struct.class, type.storageType()); + } + + @Test + void testExtensionName() { + VariantType type = VariantType.INSTANCE; + assertEquals("parquet.variant", type.extensionName()); + } + + @Test + void testExtensionEquals() { + VariantType type1 = VariantType.INSTANCE; + VariantType type2 = VariantType.INSTANCE; + + assertTrue(type1.extensionEquals(type2)); + } + + @Test + void testIsComplex() { + VariantType type = VariantType.INSTANCE; + assertFalse(type.isComplex()); + } + + @Test + void testSerialize() { + VariantType type = VariantType.INSTANCE; + String serialized = type.serialize(); + assertEquals("", serialized); + } + + @Test + void testDeserializeValid() { + VariantType type = VariantType.INSTANCE; + ArrowType storageType = ArrowType.Struct.INSTANCE; + + ArrowType deserialized = assertDoesNotThrow(() -> type.deserialize(storageType, "")); + assertInstanceOf(VariantType.class, deserialized); + assertEquals(VariantType.INSTANCE, deserialized); + } + + @Test + void testDeserializeInvalidStorageType() { + VariantType type = VariantType.INSTANCE; + ArrowType wrongStorageType = ArrowType.Utf8.INSTANCE; + + assertThrows(UnsupportedOperationException.class, () -> type.deserialize(wrongStorageType, "")); + } + + @Test + void testGetNewVector() { + VariantType type = VariantType.INSTANCE; + try (FieldVector vector = + type.getNewVector("variant_field", FieldType.nullable(type), allocator)) { + assertInstanceOf(VariantVector.class, vector); + assertEquals("variant_field", vector.getField().getName()); + assertEquals(type, vector.getField().getType()); + } + } + + @Test + void testGetNewVectorWithNullableFieldType() { + VariantType type = VariantType.INSTANCE; + FieldType nullableFieldType = FieldType.nullable(type); + + try (FieldVector vector = type.getNewVector("nullable_variant", nullableFieldType, allocator)) { + assertInstanceOf(VariantVector.class, vector); + assertEquals("nullable_variant", vector.getField().getName()); + assertTrue(vector.getField().isNullable()); + } + } + + @Test + void testGetNewVectorWithNonNullableFieldType() { + VariantType type = VariantType.INSTANCE; + FieldType nonNullableFieldType = FieldType.notNullable(type); + + try (FieldVector vector = + type.getNewVector("non_nullable_variant", nonNullableFieldType, allocator)) { + assertInstanceOf(VariantVector.class, vector); + assertEquals("non_nullable_variant", vector.getField().getName()); + } + } + + @Test + void testIpcRoundTrip() { + VariantType type = VariantType.INSTANCE; + + Schema schema = new Schema(Collections.singletonList(Field.nullable("variant", type))); + byte[] serialized = schema.serializeAsMessage(); + Schema deserialized = Schema.deserializeMessage(ByteBuffer.wrap(serialized)); + assertEquals(schema, deserialized); + } + + @Test + void testVectorIpcRoundTrip() throws IOException { + VariantType type = VariantType.INSTANCE; + + try (FieldVector vector = type.getNewVector("field", FieldType.nullable(type), allocator); + ArrowBuf metadataBuf1 = allocator.buffer(10); + ArrowBuf valueBuf1 = allocator.buffer(10); + ArrowBuf metadataBuf2 = allocator.buffer(10); + ArrowBuf valueBuf2 = allocator.buffer(10)) { + VariantVector variantVector = (VariantVector) vector; + + byte[] metadata1 = new byte[] {1, 2, 3}; + byte[] value1 = new byte[] {4, 5, 6, 7}; + metadataBuf1.setBytes(0, metadata1); + valueBuf1.setBytes(0, value1); + + byte[] metadata2 = new byte[] {8, 9}; + byte[] value2 = new byte[] {10, 11, 12}; + metadataBuf2.setBytes(0, metadata2); + valueBuf2.setBytes(0, value2); + + NullableVariantHolder holder1 = new NullableVariantHolder(); + holder1.isSet = 1; + holder1.metadataStart = 0; + holder1.metadataEnd = metadata1.length; + holder1.metadataBuffer = metadataBuf1; + holder1.valueStart = 0; + holder1.valueEnd = value1.length; + holder1.valueBuffer = valueBuf1; + + NullableVariantHolder holder2 = new NullableVariantHolder(); + holder2.isSet = 1; + holder2.metadataStart = 0; + holder2.metadataEnd = metadata2.length; + holder2.metadataBuffer = metadataBuf2; + holder2.valueStart = 0; + holder2.valueEnd = value2.length; + holder2.valueBuffer = valueBuf2; + + variantVector.setSafe(0, holder1); + variantVector.setNull(1); + variantVector.setSafe(2, holder2); + variantVector.setValueCount(3); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (VectorSchemaRoot root = new VectorSchemaRoot(Collections.singletonList(variantVector)); + ArrowStreamWriter writer = + new ArrowStreamWriter(root, new DictionaryProvider.MapDictionaryProvider(), baos)) { + writer.start(); + writer.writeBatch(); + } + + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator)) { + assertTrue(reader.loadNextBatch()); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertEquals(3, root.getRowCount()); + assertEquals( + new Schema(Collections.singletonList(variantVector.getField())), root.getSchema()); + + VariantVector actual = assertInstanceOf(VariantVector.class, root.getVector("field")); + assertFalse(actual.isNull(0)); + assertTrue(actual.isNull(1)); + assertFalse(actual.isNull(2)); + + NullableVariantHolder result1 = new NullableVariantHolder(); + actual.get(0, result1); + assertEquals(1, result1.isSet); + assertEquals(metadata1.length, result1.metadataEnd - result1.metadataStart); + assertEquals(value1.length, result1.valueEnd - result1.valueStart); + + assertNull(actual.getObject(1)); + + NullableVariantHolder result2 = new NullableVariantHolder(); + actual.get(2, result2); + assertEquals(1, result2.isSet); + assertEquals(metadata2.length, result2.metadataEnd - result2.metadataStart); + assertEquals(value2.length, result2.valueEnd - result2.valueStart); + } + } + } + + @Test + void testSingleton() { + VariantType type1 = VariantType.INSTANCE; + VariantType type2 = VariantType.INSTANCE; + + // Same instance + assertSame(type1, type2); + assertTrue(type1.extensionEquals(type2)); + } + + @Test + void testExtensionTypeRegistry() { + // VariantType should be automatically registered via static initializer + ArrowType.ExtensionType registeredType = + ExtensionTypeRegistry.lookup(VariantType.EXTENSION_NAME); + assertNotNull(registeredType); + assertInstanceOf(VariantType.class, registeredType); + assertEquals(VariantType.INSTANCE, registeredType); + } + + @Test + void testFieldMetadata() { + Map metadata = new HashMap<>(); + metadata.put("key1", "value1"); + metadata.put("key2", "value2"); + + FieldType fieldType = new FieldType(true, VariantType.INSTANCE, null, metadata); + try (VariantVector vector = new VariantVector("test", allocator)) { + Field field = new Field("test", fieldType, VariantVector.createVariantChildFields()); + + // Field metadata includes both custom metadata and extension type metadata + Map fieldMetadata = field.getMetadata(); + assertEquals("value1", fieldMetadata.get("key1")); + assertEquals("value2", fieldMetadata.get("key2")); + // Extension type metadata is also present + assertTrue(fieldMetadata.containsKey("ARROW:extension:name")); + assertTrue(fieldMetadata.containsKey("ARROW:extension:metadata")); + } + } + + @Test + void testFieldChildren() { + try (VariantVector vector = new VariantVector("test", allocator)) { + Field field = vector.getField(); + + assertNotNull(field.getChildren()); + assertEquals(2, field.getChildren().size()); + + Field metadataField = field.getChildren().get(0); + assertEquals(VariantVector.METADATA_VECTOR_NAME, metadataField.getName()); + assertEquals(ArrowType.Binary.INSTANCE, metadataField.getType()); + + Field valueField = field.getChildren().get(1); + assertEquals(VariantVector.VALUE_VECTOR_NAME, valueField.getName()); + assertEquals(ArrowType.Binary.INSTANCE, valueField.getType()); + } + } +} diff --git a/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java new file mode 100644 index 0000000000..1c172e304f --- /dev/null +++ b/arrow-variant/src/test/java/org/apache/arrow/variant/extension/TestVariantVector.java @@ -0,0 +1,844 @@ +/* + * 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.arrow.variant.extension; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +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 org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.variant.Variant; +import org.apache.arrow.variant.holders.NullableVariantHolder; +import org.apache.arrow.variant.holders.VariantHolder; +import org.apache.arrow.variant.impl.VariantReaderImpl; +import org.apache.arrow.variant.impl.VariantWriterImpl; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for VariantVector, VariantWriterImpl, and VariantReaderImpl. */ +class TestVariantVector { + + private BufferAllocator allocator; + + @BeforeEach + void beforeEach() { + allocator = new RootAllocator(); + } + + @AfterEach + void afterEach() { + allocator.close(); + } + + private VariantHolder createHolder( + ArrowBuf metadataBuf, byte[] metadata, ArrowBuf valueBuf, byte[] value) { + VariantHolder holder = new VariantHolder(); + holder.metadataStart = 0; + holder.metadataEnd = metadata.length; + holder.metadataBuffer = metadataBuf; + holder.valueStart = 0; + holder.valueEnd = value.length; + holder.valueBuffer = valueBuf; + return holder; + } + + private NullableVariantHolder createNullableHolder( + ArrowBuf metadataBuf, byte[] metadata, ArrowBuf valueBuf, byte[] value) { + NullableVariantHolder holder = new NullableVariantHolder(); + holder.isSet = 1; + holder.metadataStart = 0; + holder.metadataEnd = metadata.length; + holder.metadataBuffer = metadataBuf; + holder.valueStart = 0; + holder.valueEnd = value.length; + holder.valueBuffer = valueBuf; + return holder; + } + + private NullableVariantHolder createNullHolder() { + NullableVariantHolder holder = new NullableVariantHolder(); + holder.isSet = 0; + return holder; + } + + // ========== Basic Vector Tests ========== + + @Test + void testVectorCreation() { + try (VariantVector vector = new VariantVector("test", allocator)) { + assertNotNull(vector); + assertEquals("test", vector.getField().getName()); + assertNotNull(vector.getMetadataVector()); + assertNotNull(vector.getValueVector()); + } + } + + @Test + void testSetAndGet() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5, 6, 7}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + // Retrieve and verify + NullableVariantHolder result = new NullableVariantHolder(); + vector.get(0, result); + + assertEquals(1, result.isSet); + assertEquals(metadata.length, result.metadataEnd - result.metadataStart); + assertEquals(value.length, result.valueEnd - result.valueStart); + + byte[] actualMetadata = new byte[metadata.length]; + byte[] actualValue = new byte[value.length]; + result.metadataBuffer.getBytes(result.metadataStart, actualMetadata); + result.valueBuffer.getBytes(result.valueStart, actualValue); + + assertArrayEquals(metadata, actualMetadata); + assertArrayEquals(value, actualValue); + } + } + + @Test + void testSetNull() { + try (VariantVector vector = new VariantVector("test", allocator)) { + NullableVariantHolder holder = createNullHolder(); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + assertTrue(vector.isNull(0)); + + NullableVariantHolder result = new NullableVariantHolder(); + vector.get(0, result); + assertEquals(0, result.isSet); + } + } + + @Test + void testMultipleValues() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf1 = allocator.buffer(10); + ArrowBuf valueBuf1 = allocator.buffer(10); + ArrowBuf metadataBuf2 = allocator.buffer(10); + ArrowBuf valueBuf2 = allocator.buffer(10)) { + + byte[] metadata1 = new byte[] {1, 2}; + byte[] value1 = new byte[] {3, 4, 5}; + metadataBuf1.setBytes(0, metadata1); + valueBuf1.setBytes(0, value1); + + NullableVariantHolder holder1 = + createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1); + + byte[] metadata2 = new byte[] {6, 7, 8}; + byte[] value2 = new byte[] {9, 10}; + metadataBuf2.setBytes(0, metadata2); + valueBuf2.setBytes(0, value2); + + NullableVariantHolder holder2 = + createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2); + + vector.setSafe(0, holder1); + vector.setSafe(1, holder2); + vector.setValueCount(2); + + // Verify first value + NullableVariantHolder result1 = new NullableVariantHolder(); + vector.get(0, result1); + assertEquals(1, result1.isSet); + + byte[] actualMetadata1 = new byte[metadata1.length]; + byte[] actualValue1 = new byte[value1.length]; + result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1); + result1.valueBuffer.getBytes(result1.valueStart, actualValue1); + assertArrayEquals(metadata1, actualMetadata1); + assertArrayEquals(value1, actualValue1); + + // Verify second value + NullableVariantHolder result2 = new NullableVariantHolder(); + vector.get(1, result2); + assertEquals(1, result2.isSet); + + byte[] actualMetadata2 = new byte[metadata2.length]; + byte[] actualValue2 = new byte[value2.length]; + result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2); + result2.valueBuffer.getBytes(result2.valueStart, actualValue2); + assertArrayEquals(metadata2, actualMetadata2); + assertArrayEquals(value2, actualValue2); + } + } + + @Test + void testNonNullableHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5, 6}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + VariantHolder holder = createHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + assertFalse(vector.isNull(0)); + + NullableVariantHolder result = new NullableVariantHolder(); + vector.get(0, result); + assertEquals(1, result.isSet); + } + } + + // ========== Writer Tests ========== + + @Test + void testWriteWithVariantHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + VariantWriterImpl writer = new VariantWriterImpl(vector); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2}; + byte[] value = new byte[] {3, 4, 5}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + VariantHolder holder = createHolder(metadataBuf, metadata, valueBuf, value); + + writer.setPosition(0); + writer.write(holder); + + assertEquals(1, vector.getValueCount()); + assertFalse(vector.isNull(0)); + } + } + + @Test + void testWriteWithNullableVariantHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + VariantWriterImpl writer = new VariantWriterImpl(vector); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2}; + byte[] value = new byte[] {3, 4, 5}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + writer.setPosition(0); + writer.write(holder); + + assertEquals(1, vector.getValueCount()); + assertFalse(vector.isNull(0)); + } + } + + @Test + void testWriteWithNullableVariantHolderNull() { + try (VariantVector vector = new VariantVector("test", allocator); + VariantWriterImpl writer = new VariantWriterImpl(vector)) { + + NullableVariantHolder holder = createNullHolder(); + + writer.setPosition(0); + writer.write(holder); + + assertEquals(1, vector.getValueCount()); + assertTrue(vector.isNull(0)); + } + } + + @Test + void testWriteExtensionWithUnsupportedType() { + try (VariantVector vector = new VariantVector("test", allocator); + VariantWriterImpl writer = new VariantWriterImpl(vector)) { + + writer.setPosition(0); + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> writer.writeExtension("invalid-type")); + + assertTrue(exception.getMessage().contains("Unsupported type for Variant")); + } + } + + @Test + void testWriteWithUnsupportedHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + VariantWriterImpl writer = new VariantWriterImpl(vector)) { + + ExtensionHolder unsupportedHolder = + new ExtensionHolder() { + @Override + public ArrowType type() { + return VariantType.INSTANCE; + } + }; + + writer.setPosition(0); + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> writer.write(unsupportedHolder)); + + assertTrue(exception.getMessage().contains("Unsupported type for Variant")); + } + } + + // ========== Reader Tests ========== + + @Test + void testReaderReadWithNullableVariantHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5, 6}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + reader.setPosition(0); + + NullableVariantHolder result = new NullableVariantHolder(); + reader.read(result); + + assertEquals(1, result.isSet); + assertEquals(metadata.length, result.metadataEnd - result.metadataStart); + assertEquals(value.length, result.valueEnd - result.valueStart); + } + } + + @Test + void testReaderReadWithNullableVariantHolderNull() { + try (VariantVector vector = new VariantVector("test", allocator)) { + vector.setNull(0); + vector.setValueCount(1); + + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + reader.setPosition(0); + + NullableVariantHolder holder = new NullableVariantHolder(); + reader.read(holder); + + assertEquals(0, holder.isSet); + } + } + + @Test + void testReaderIsSet() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1}; + byte[] value = new byte[] {2}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setNull(1); + vector.setValueCount(2); + + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + + reader.setPosition(0); + assertTrue(reader.isSet()); + + reader.setPosition(1); + assertFalse(reader.isSet()); + } + } + + @Test + void testReaderGetMinorType() { + try (VariantVector vector = new VariantVector("test", allocator)) { + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + assertEquals(vector.getMinorType(), reader.getMinorType()); + } + } + + @Test + void testReaderGetField() { + try (VariantVector vector = new VariantVector("test", allocator)) { + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + assertEquals(vector.getField(), reader.getField()); + assertEquals("test", reader.getField().getName()); + } + } + + @Test + void testReaderReadWithNonNullableVariantHolder() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5, 6}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + VariantReaderImpl reader = (VariantReaderImpl) vector.getReader(); + reader.setPosition(0); + + VariantHolder result = new VariantHolder(); + reader.read(result); + + // Verify the data was read correctly + byte[] actualMetadata = new byte[metadata.length]; + byte[] actualValue = new byte[value.length]; + result.metadataBuffer.getBytes(result.metadataStart, actualMetadata); + result.valueBuffer.getBytes(result.valueStart, actualValue); + + assertArrayEquals(metadata, actualMetadata); + assertArrayEquals(value, actualValue); + assertEquals(1, result.isSet); + } + } + + // ========== Transfer Pair Tests ========== + + @Test + void testTransferPair() { + try (VariantVector fromVector = new VariantVector("from", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5, 6, 7}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + fromVector.setSafe(0, holder); + fromVector.setValueCount(1); + + org.apache.arrow.vector.util.TransferPair transferPair = + fromVector.getTransferPair(allocator); + VariantVector toVector = (VariantVector) transferPair.getTo(); + + transferPair.transfer(); + + assertEquals(0, fromVector.getValueCount()); + assertEquals(1, toVector.getValueCount()); + + NullableVariantHolder result = new NullableVariantHolder(); + toVector.get(0, result); + assertEquals(1, result.isSet); + + byte[] actualMetadata = new byte[metadata.length]; + byte[] actualValue = new byte[value.length]; + result.metadataBuffer.getBytes(result.metadataStart, actualMetadata); + result.valueBuffer.getBytes(result.valueStart, actualValue); + + assertArrayEquals(metadata, actualMetadata); + assertArrayEquals(value, actualValue); + + toVector.close(); + } + } + + @Test + void testSplitAndTransfer() { + try (VariantVector fromVector = new VariantVector("from", allocator); + ArrowBuf metadataBuf1 = allocator.buffer(10); + ArrowBuf valueBuf1 = allocator.buffer(10); + ArrowBuf metadataBuf2 = allocator.buffer(10); + ArrowBuf valueBuf2 = allocator.buffer(10); + ArrowBuf metadataBuf3 = allocator.buffer(10); + ArrowBuf valueBuf3 = allocator.buffer(10)) { + + byte[] metadata1 = new byte[] {1}; + byte[] value1 = new byte[] {2, 3}; + metadataBuf1.setBytes(0, metadata1); + valueBuf1.setBytes(0, value1); + + byte[] metadata2 = new byte[] {4, 5}; + byte[] value2 = new byte[] {6}; + metadataBuf2.setBytes(0, metadata2); + valueBuf2.setBytes(0, value2); + + byte[] metadata3 = new byte[] {7, 8, 9}; + byte[] value3 = new byte[] {10, 11, 12}; + metadataBuf3.setBytes(0, metadata3); + valueBuf3.setBytes(0, value3); + + NullableVariantHolder holder1 = + createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1); + NullableVariantHolder holder2 = + createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2); + NullableVariantHolder holder3 = + createNullableHolder(metadataBuf3, metadata3, valueBuf3, value3); + + fromVector.setSafe(0, holder1); + fromVector.setSafe(1, holder2); + fromVector.setSafe(2, holder3); + fromVector.setValueCount(3); + + org.apache.arrow.vector.util.TransferPair transferPair = + fromVector.getTransferPair(allocator); + VariantVector toVector = (VariantVector) transferPair.getTo(); + + // Split and transfer indices 1-2 (middle and last) + transferPair.splitAndTransfer(1, 2); + + assertEquals(2, toVector.getValueCount()); + + // Verify transferred values + NullableVariantHolder result1 = new NullableVariantHolder(); + toVector.get(0, result1); + assertEquals(1, result1.isSet); + + byte[] actualMetadata1 = new byte[metadata2.length]; + byte[] actualValue1 = new byte[value2.length]; + result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1); + result1.valueBuffer.getBytes(result1.valueStart, actualValue1); + assertArrayEquals(metadata2, actualMetadata1); + assertArrayEquals(value2, actualValue1); + + NullableVariantHolder result2 = new NullableVariantHolder(); + toVector.get(1, result2); + assertEquals(1, result2.isSet); + + byte[] actualMetadata2 = new byte[metadata3.length]; + byte[] actualValue2 = new byte[value3.length]; + result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2); + result2.valueBuffer.getBytes(result2.valueStart, actualValue2); + assertArrayEquals(metadata3, actualMetadata2); + assertArrayEquals(value3, actualValue2); + + toVector.close(); + } + } + + @Test + void testCopyValueSafe() { + try (VariantVector fromVector = new VariantVector("from", allocator); + VariantVector toVector = new VariantVector("to", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2}; + byte[] value = new byte[] {3, 4, 5}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + fromVector.setSafe(0, holder); + fromVector.setValueCount(1); + + org.apache.arrow.vector.util.TransferPair transferPair = + fromVector.makeTransferPair(toVector); + + transferPair.copyValueSafe(0, 0); + toVector.setValueCount(1); + + // Verify the value was copied + NullableVariantHolder result = new NullableVariantHolder(); + toVector.get(0, result); + assertEquals(1, result.isSet); + + byte[] actualMetadata = new byte[metadata.length]; + byte[] actualValue = new byte[value.length]; + result.metadataBuffer.getBytes(result.metadataStart, actualMetadata); + result.valueBuffer.getBytes(result.valueStart, actualValue); + + assertArrayEquals(metadata, actualMetadata); + assertArrayEquals(value, actualValue); + + // Original vector should still have the value + NullableVariantHolder originalResult = new NullableVariantHolder(); + fromVector.get(0, originalResult); + assertEquals(1, originalResult.isSet); + } + } + + @Test + void testGetTransferPairWithField() { + try (VariantVector fromVector = new VariantVector("from", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1}; + byte[] value = new byte[] {2}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + fromVector.setSafe(0, holder); + fromVector.setValueCount(1); + + org.apache.arrow.vector.util.TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField(), allocator); + VariantVector toVector = (VariantVector) transferPair.getTo(); + + transferPair.transfer(); + + assertEquals(1, toVector.getValueCount()); + assertEquals(fromVector.getField().getName(), toVector.getField().getName()); + + toVector.close(); + } + } + + // ========== Copy Operations Tests ========== + + @Test + void testCopyFrom() { + try (VariantVector fromVector = new VariantVector("from", allocator); + VariantVector toVector = new VariantVector("to", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2, 3}; + byte[] value = new byte[] {4, 5}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + fromVector.setSafe(0, holder); + fromVector.setValueCount(1); + + toVector.allocateNew(); + toVector.copyFrom(0, 0, fromVector); + toVector.setValueCount(1); + + NullableVariantHolder result = new NullableVariantHolder(); + toVector.get(0, result); + assertEquals(1, result.isSet); + + byte[] actualMetadata = new byte[metadata.length]; + byte[] actualValue = new byte[value.length]; + result.metadataBuffer.getBytes(result.metadataStart, actualMetadata); + result.valueBuffer.getBytes(result.valueStart, actualValue); + + assertArrayEquals(metadata, actualMetadata); + assertArrayEquals(value, actualValue); + } + } + + @Test + void testCopyFromSafe() { + try (VariantVector fromVector = new VariantVector("from", allocator); + VariantVector toVector = new VariantVector("to", allocator); + ArrowBuf metadataBuf1 = allocator.buffer(10); + ArrowBuf valueBuf1 = allocator.buffer(10); + ArrowBuf metadataBuf2 = allocator.buffer(10); + ArrowBuf valueBuf2 = allocator.buffer(10)) { + + byte[] metadata1 = new byte[] {1}; + byte[] value1 = new byte[] {2, 3}; + metadataBuf1.setBytes(0, metadata1); + valueBuf1.setBytes(0, value1); + + NullableVariantHolder holder1 = + createNullableHolder(metadataBuf1, metadata1, valueBuf1, value1); + + byte[] metadata2 = new byte[] {4, 5}; + byte[] value2 = new byte[] {6}; + metadataBuf2.setBytes(0, metadata2); + valueBuf2.setBytes(0, value2); + + NullableVariantHolder holder2 = + createNullableHolder(metadataBuf2, metadata2, valueBuf2, value2); + + fromVector.setSafe(0, holder1); + fromVector.setSafe(1, holder2); + fromVector.setValueCount(2); + + // Copy without pre-allocating toVector + for (int i = 0; i < 2; i++) { + toVector.copyFromSafe(i, i, fromVector); + } + toVector.setValueCount(2); + + // Verify both values + NullableVariantHolder result1 = new NullableVariantHolder(); + toVector.get(0, result1); + assertEquals(1, result1.isSet); + + byte[] actualMetadata1 = new byte[metadata1.length]; + byte[] actualValue1 = new byte[value1.length]; + result1.metadataBuffer.getBytes(result1.metadataStart, actualMetadata1); + result1.valueBuffer.getBytes(result1.valueStart, actualValue1); + assertArrayEquals(metadata1, actualMetadata1); + assertArrayEquals(value1, actualValue1); + + NullableVariantHolder result2 = new NullableVariantHolder(); + toVector.get(1, result2); + assertEquals(1, result2.isSet); + + byte[] actualMetadata2 = new byte[metadata2.length]; + byte[] actualValue2 = new byte[value2.length]; + result2.metadataBuffer.getBytes(result2.metadataStart, actualMetadata2); + result2.valueBuffer.getBytes(result2.valueStart, actualValue2); + assertArrayEquals(metadata2, actualMetadata2); + assertArrayEquals(value2, actualValue2); + } + } + + @Test + void testCopyFromWithNulls() { + try (VariantVector fromVector = new VariantVector("from", allocator); + VariantVector toVector = new VariantVector("to", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1}; + byte[] value = new byte[] {2}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + fromVector.setSafe(0, holder); + fromVector.setNull(1); + fromVector.setSafe(2, holder); + fromVector.setValueCount(3); + + toVector.allocateNew(); + for (int i = 0; i < 3; i++) { + toVector.copyFromSafe(i, i, fromVector); + } + toVector.setValueCount(3); + + assertFalse(toVector.isNull(0)); + assertTrue(toVector.isNull(1)); + assertFalse(toVector.isNull(2)); + } + } + + // ========== GetObject Tests ========== + + @Test + void testGetObject() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1, 2}; + byte[] value = new byte[] {3, 4, 5}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + Object obj = vector.getObject(0); + assertNotNull(obj); + assertTrue(obj instanceof Variant); + assertEquals(new Variant(metadata, value), obj); + } + } + + @Test + void testGetObjectNull() { + try (VariantVector vector = new VariantVector("test", allocator)) { + vector.setNull(0); + vector.setValueCount(1); + + Object obj = vector.getObject(0); + assertNull(obj); + } + } + + // ========== Allocate and Capacity Tests ========== + + @Test + void testAllocateNew() { + try (VariantVector vector = new VariantVector("test", allocator)) { + vector.allocateNew(); + assertTrue(vector.getValueCapacity() > 0); + } + } + + @Test + void testSetInitialCapacity() { + try (VariantVector vector = new VariantVector("test", allocator)) { + vector.setInitialCapacity(100); + vector.allocateNew(); + assertTrue(vector.getValueCapacity() >= 100); + } + } + + @Test + void testClearAndReuse() { + try (VariantVector vector = new VariantVector("test", allocator); + ArrowBuf metadataBuf = allocator.buffer(10); + ArrowBuf valueBuf = allocator.buffer(10)) { + + byte[] metadata = new byte[] {1}; + byte[] value = new byte[] {2}; + metadataBuf.setBytes(0, metadata); + valueBuf.setBytes(0, value); + + NullableVariantHolder holder = createNullableHolder(metadataBuf, metadata, valueBuf, value); + + vector.setSafe(0, holder); + vector.setValueCount(1); + + assertFalse(vector.isNull(0)); + + vector.clear(); + vector.allocateNew(); + + // After clear, vector should be empty + assertEquals(0, vector.getValueCount()); + } + } +} diff --git a/bom/pom.xml b/bom/pom.xml index 171d0bc5e9..2d1085b160 100644 --- a/bom/pom.xml +++ b/bom/pom.xml @@ -23,13 +23,13 @@ under the License. org.apache apache - 33 - + 38 + org.apache.arrow arrow-bom - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT pom Arrow Bill of Materials @@ -78,12 +78,12 @@ under the License. - + - 11 - 11 - 11 - 11 + 17 + 17 + 17 + 17 @@ -165,7 +165,7 @@ under the License. ${project.version} - org.apache.arrow + org.apache.arrow.gandiva arrow-gandiva ${project.version} @@ -194,6 +194,11 @@ under the License. arrow-tools ${project.version} + + org.apache.arrow + arrow-variant + ${project.version} + @@ -203,12 +208,12 @@ under the License. com.diffplug.spotless spotless-maven-plugin - 2.44.3 + 3.8.0 org.codehaus.mojo versions-maven-plugin - 2.18.0 + 2.21.0 @@ -230,7 +235,7 @@ under the License. ${maven.multiModuleProjectDirectory}/dev/license/asf-xml.license (<configuration|<project) - + diff --git a/c/pom.xml b/c/pom.xml index c90b6dc0ef..27b6619c4c 100644 --- a/c/pom.xml +++ b/c/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-c-data diff --git a/c/src/main/cpp/jni_wrapper.cc b/c/src/main/cpp/jni_wrapper.cc index 35c2b7787e..3d7a194563 100644 --- a/c/src/main/cpp/jni_wrapper.cc +++ b/c/src/main/cpp/jni_wrapper.cc @@ -205,8 +205,9 @@ void TryCopyLastError(JNIEnv* env, InnerPrivateData* private_data) { return; } + jsize error_bytes_len = env->GetArrayLength(arr); char* error_str = reinterpret_cast(error_bytes); - private_data->last_error_ = std::string(error_str, std::strlen(error_str)); + private_data->last_error_ = std::string(error_str, error_bytes_len); env->ReleaseByteArrayElements(arr, error_bytes, JNI_ABORT); } @@ -326,19 +327,20 @@ void ArrowArrayStreamRelease(ArrowArrayStream* stream) { jint JNI_OnLoad(JavaVM* vm, void* reserved) { JNIEnv* env; - if (vm->GetEnv(reinterpret_cast(&env), JNI_VERSION) != JNI_OK) { - return JNI_ERR; + const int err_code = vm->GetEnv(reinterpret_cast(&env), JNI_VERSION); + if (err_code != JNI_OK) { + return err_code; } JNI_METHOD_START - kObjectClass = CreateGlobalClassReference(env, "Ljava/lang/Object;"); + kObjectClass = CreateGlobalClassReference(env, "java/lang/Object"); kRuntimeExceptionClass = - CreateGlobalClassReference(env, "Ljava/lang/RuntimeException;"); + CreateGlobalClassReference(env, "java/lang/RuntimeException"); kPrivateDataClass = - CreateGlobalClassReference(env, "Lorg/apache/arrow/c/jni/PrivateData;"); + CreateGlobalClassReference(env, "org/apache/arrow/c/jni/PrivateData"); kCDataExceptionClass = - CreateGlobalClassReference(env, "Lorg/apache/arrow/c/jni/CDataJniException;"); + CreateGlobalClassReference(env, "org/apache/arrow/c/jni/CDataJniException"); kStreamPrivateDataClass = CreateGlobalClassReference( - env, "Lorg/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData;"); + env, "org/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData"); kPrivateDataLastErrorField = GetFieldID(env, kStreamPrivateDataClass, "lastError", "[B"); diff --git a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java index b74fb1b473..f31a8a1faa 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrayImporter.java +++ b/c/src/main/java/org/apache/arrow/c/ArrayImporter.java @@ -58,7 +58,6 @@ void importArray(ArrowArray src) { ArrowArray ownedArray = ArrowArray.allocateNew(allocator); ownedArray.save(snapshot); src.markReleased(); - src.close(); recursionLevel = 0; diff --git a/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java b/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java index 07a88cd8d7..34a9c4ec03 100644 --- a/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java +++ b/c/src/main/java/org/apache/arrow/c/ArrowArrayStreamReader.java @@ -44,7 +44,6 @@ final class ArrowArrayStreamReader extends ArrowReader { this.ownedStream = ArrowArrayStream.allocateNew(allocator); this.ownedStream.save(snapshot); stream.markReleased(); - stream.close(); } @Override diff --git a/c/src/main/java/org/apache/arrow/c/Data.java b/c/src/main/java/org/apache/arrow/c/Data.java index 0b4da33b4e..f9d2ee4542 100644 --- a/c/src/main/java/org/apache/arrow/c/Data.java +++ b/c/src/main/java/org/apache/arrow/c/Data.java @@ -231,6 +231,22 @@ public static void exportArrayStream( new ArrayStreamExporter(allocator).export(out, reader); } + /** + * Equivalent to calling {@link #importField(BufferAllocator, ArrowSchema, + * CDataDictionaryProvider, boolean) importField(allocator, schema, provider, true)}. + * + * @param allocator Buffer allocator for allocating dictionary vectors + * @param schema C data interface struct representing the field [inout] + * @param provider A dictionary provider will be initialized with empty dictionary vectors + * (optional) + * @return Imported field object + * @see #importField(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean) + */ + public static Field importField( + BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { + return importField(allocator, schema, provider, true); + } + /** * Import Java Field from the C data interface. * @@ -241,19 +257,42 @@ public static void exportArrayStream( * @param schema C data interface struct representing the field [inout] * @param provider A dictionary provider will be initialized with empty dictionary vectors * (optional) + * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method + * completes. * @return Imported field object */ public static Field importField( - BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { + BufferAllocator allocator, + ArrowSchema schema, + CDataDictionaryProvider provider, + boolean closeImportedStructs) { try { SchemaImporter importer = new SchemaImporter(allocator); return importer.importField(schema, provider); } finally { schema.release(); - schema.close(); + if (closeImportedStructs) { + schema.close(); + } } } + /** + * Equivalent to calling {@link #importSchema(BufferAllocator, ArrowSchema, + * CDataDictionaryProvider, boolean) importSchema(allocator, schema, provider, true)}. + * + * @param allocator Buffer allocator for allocating dictionary vectors + * @param schema C data interface struct representing the field + * @param provider A dictionary provider will be initialized with empty dictionary vectors + * (optional) + * @return Imported schema object + * @see #importSchema(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean) + */ + public static Schema importSchema( + BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { + return importSchema(allocator, schema, provider, true); + } + /** * Import Java Schema from the C data interface. * @@ -264,11 +303,16 @@ public static Field importField( * @param schema C data interface struct representing the field * @param provider A dictionary provider will be initialized with empty dictionary vectors * (optional) + * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method + * completes. * @return Imported schema object */ public static Schema importSchema( - BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { - Field structField = importField(allocator, schema, provider); + BufferAllocator allocator, + ArrowSchema schema, + CDataDictionaryProvider provider, + boolean closeImportedStructs) { + Field structField = importField(allocator, schema, provider, closeImportedStructs); if (structField.getType().getTypeID() != ArrowTypeID.Struct) { throw new IllegalArgumentException( "Cannot import schema: ArrowSchema describes non-struct type"); @@ -276,24 +320,67 @@ public static Schema importSchema( return new Schema(structField.getChildren(), structField.getMetadata()); } + /** + * Equivalent to calling {@link #importIntoVector(BufferAllocator, ArrowArray, FieldVector, + * DictionaryProvider, boolean)} importIntoVector(allocator, array, vector, provider, true)}. + * + * @param allocator Buffer allocator + * @param array C data interface struct holding the array data + * @param vector Imported vector object [out] + * @param provider Dictionary provider to load dictionary vectors to (optional) + * @see #importIntoVector(BufferAllocator, ArrowArray, FieldVector, DictionaryProvider, boolean) + */ + public static void importIntoVector( + BufferAllocator allocator, + ArrowArray array, + FieldVector vector, + DictionaryProvider provider) { + importIntoVector(allocator, array, vector, provider, true); + } + /** * Import Java vector from the C data interface. * - *

The ArrowArray struct has its contents moved (as per the C data interface specification) to - * a private object held alive by the resulting array. + *

On successful completion, the ArrowArray struct will have been moved (as per the C data + * interface specification) to a private object held alive by the resulting array. * * @param allocator Buffer allocator * @param array C data interface struct holding the array data * @param vector Imported vector object [out] * @param provider Dictionary provider to load dictionary vectors to (optional) + * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method + * completes successfully. */ public static void importIntoVector( BufferAllocator allocator, ArrowArray array, FieldVector vector, - DictionaryProvider provider) { + DictionaryProvider provider, + boolean closeImportedStructs) { ArrayImporter importer = new ArrayImporter(allocator, vector, provider); importer.importArray(array); + if (closeImportedStructs) { + array.close(); + } + } + + /** + * Equivalent to calling {@link #importVector(BufferAllocator, ArrowArray, ArrowSchema, + * CDataDictionaryProvider, boolean) importVector(allocator, array, schema, provider, true)}. + * + * @param allocator Buffer allocator for allocating the output FieldVector + * @param array C data interface struct holding the array data + * @param schema C data interface struct holding the array type + * @param provider Dictionary provider to load dictionary vectors to (optional) + * @return Imported vector object + * @see #importVector(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider, boolean) + */ + public static FieldVector importVector( + BufferAllocator allocator, + ArrowArray array, + ArrowSchema schema, + CDataDictionaryProvider provider) { + return importVector(allocator, array, schema, provider, true); } /** @@ -307,19 +394,42 @@ public static void importIntoVector( * @param array C data interface struct holding the array data * @param schema C data interface struct holding the array type * @param provider Dictionary provider to load dictionary vectors to (optional) + * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method + * completes successfully and the ArrowSchema struct will be always be closed. * @return Imported vector object */ public static FieldVector importVector( BufferAllocator allocator, ArrowArray array, ArrowSchema schema, - CDataDictionaryProvider provider) { - Field field = importField(allocator, schema, provider); + CDataDictionaryProvider provider, + boolean closeImportedStructs) { + Field field = importField(allocator, schema, provider, closeImportedStructs); FieldVector vector = field.createVector(allocator); - importIntoVector(allocator, array, vector, provider); + importIntoVector(allocator, array, vector, provider, closeImportedStructs); return vector; } + /** + * Equivalent to calling {@link #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray, + * VectorSchemaRoot, DictionaryProvider, boolean) importIntoVectorSchemaRoot(allocator, array, + * root, provider, true)}. + * + * @param allocator Buffer allocator + * @param array C data interface struct holding the record batch data + * @param root vector schema root to load into + * @param provider Dictionary provider to load dictionary vectors to (optional) + * @see #importIntoVectorSchemaRoot(BufferAllocator, ArrowArray, VectorSchemaRoot, + * DictionaryProvider, boolean) + */ + public static void importIntoVectorSchemaRoot( + BufferAllocator allocator, + ArrowArray array, + VectorSchemaRoot root, + DictionaryProvider provider) { + importIntoVectorSchemaRoot(allocator, array, root, provider, true); + } + /** * Import record batch from the C data interface into vector schema root. * @@ -333,15 +443,18 @@ public static FieldVector importVector( * @param array C data interface struct holding the record batch data * @param root vector schema root to load into * @param provider Dictionary provider to load dictionary vectors to (optional) + * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method + * completes successfully */ public static void importIntoVectorSchemaRoot( BufferAllocator allocator, ArrowArray array, VectorSchemaRoot root, - DictionaryProvider provider) { + DictionaryProvider provider, + boolean closeImportedStructs) { try (StructVector structVector = StructVector.emptyWithDuplicates("", allocator)) { structVector.initializeChildrenFromFields(root.getSchema().getFields()); - importIntoVector(allocator, array, structVector, provider); + importIntoVector(allocator, array, structVector, provider, closeImportedStructs); StructVectorUnloader unloader = new StructVectorUnloader(structVector); VectorLoader loader = new VectorLoader(root); try (ArrowRecordBatch recordBatch = unloader.getRecordBatch()) { @@ -350,6 +463,21 @@ public static void importIntoVectorSchemaRoot( } } + /** + * Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowSchema, + * CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, schema, provider, true)}. + * + * @param allocator Buffer allocator for allocating the output VectorSchemaRoot + * @param schema C data interface struct holding the record batch schema + * @param provider Dictionary provider to load dictionary vectors to (optional) + * @return Imported vector schema root + * @see #importVectorSchemaRoot(BufferAllocator, ArrowSchema, CDataDictionaryProvider, boolean) + */ + public static VectorSchemaRoot importVectorSchemaRoot( + BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { + return importVectorSchemaRoot(allocator, schema, provider, true); + } + /** * Import Java vector schema root from a C data interface Schema. * @@ -360,11 +488,37 @@ public static void importIntoVectorSchemaRoot( * @param allocator Buffer allocator for allocating the output VectorSchemaRoot * @param schema C data interface struct holding the record batch schema * @param provider Dictionary provider to load dictionary vectors to (optional) + * @param closeImportedStructs if true, the ArrowSchema struct will be closed when this method + * completes * @return Imported vector schema root */ public static VectorSchemaRoot importVectorSchemaRoot( - BufferAllocator allocator, ArrowSchema schema, CDataDictionaryProvider provider) { - return importVectorSchemaRoot(allocator, null, schema, provider); + BufferAllocator allocator, + ArrowSchema schema, + CDataDictionaryProvider provider, + boolean closeImportedStructs) { + return importVectorSchemaRoot(allocator, null, schema, provider, closeImportedStructs); + } + + /** + * Equivalent to calling {@link #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema, + * CDataDictionaryProvider, boolean) importVectorSchemaRoot(allocator, array, schema, provider, + * true)}. + * + * @param allocator Buffer allocator for allocating the output VectorSchemaRoot + * @param array C data interface struct holding the record batch data (optional) + * @param schema C data interface struct holding the record batch schema + * @param provider Dictionary provider to load dictionary vectors to (optional) + * @return Imported vector schema root + * @see #importVectorSchemaRoot(BufferAllocator, ArrowArray, ArrowSchema, CDataDictionaryProvider, + * boolean) + */ + public static VectorSchemaRoot importVectorSchemaRoot( + BufferAllocator allocator, + ArrowArray array, + ArrowSchema schema, + CDataDictionaryProvider provider) { + return importVectorSchemaRoot(allocator, array, schema, provider, true); } /** @@ -383,29 +537,56 @@ public static VectorSchemaRoot importVectorSchemaRoot( * @param array C data interface struct holding the record batch data (optional) * @param schema C data interface struct holding the record batch schema * @param provider Dictionary provider to load dictionary vectors to (optional) + * @param closeImportedStructs if true, the ArrowArray struct will be closed when this method + * completes successfully and the ArrowSchema struct will be always be closed. * @return Imported vector schema root */ public static VectorSchemaRoot importVectorSchemaRoot( BufferAllocator allocator, ArrowArray array, ArrowSchema schema, - CDataDictionaryProvider provider) { + CDataDictionaryProvider provider, + boolean closeImportedStructs) { VectorSchemaRoot vsr = - VectorSchemaRoot.create(importSchema(allocator, schema, provider), allocator); + VectorSchemaRoot.create( + importSchema(allocator, schema, provider, closeImportedStructs), allocator); if (array != null) { - importIntoVectorSchemaRoot(allocator, array, vsr, provider); + importIntoVectorSchemaRoot(allocator, array, vsr, provider, closeImportedStructs); } return vsr; } /** - * Import an ArrowArrayStream as an {@link ArrowReader}. + * Equivalent to calling {@link #importArrayStream(BufferAllocator, ArrowArrayStream, boolean) + * importArrayStream(allocator, stream, true)}. * * @param allocator Buffer allocator for allocating the output data. * @param stream C stream interface struct to import. * @return Imported reader + * @see #importArrayStream(BufferAllocator, ArrowArrayStream, boolean) */ public static ArrowReader importArrayStream(BufferAllocator allocator, ArrowArrayStream stream) { - return new ArrowArrayStreamReader(allocator, stream); + return importArrayStream(allocator, stream, true); + } + + /** + * Import an ArrowArrayStream as an {@link ArrowReader}. + * + *

On successful completion, the ArrowArrayStream struct will have been moved (as per the C + * data interface specification) to a private object held alive by the resulting ArrowReader. + * + * @param allocator Buffer allocator for allocating the output data. + * @param stream C stream interface struct to import. + * @param closeImportedStructs if true, the ArrowArrayStream struct will be closed when this + * method completes successfully + * @return Imported reader + */ + public static ArrowReader importArrayStream( + BufferAllocator allocator, ArrowArrayStream stream, boolean closeImportedStructs) { + ArrowArrayStreamReader reader = new ArrowArrayStreamReader(allocator, stream); + if (closeImportedStructs) { + stream.close(); + } + return reader; } } diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java index cf50f9417b..f51fb25105 100644 --- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java +++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java @@ -64,13 +64,18 @@ void release() { */ ArrowBuf unsafeAssociateAllocation( BufferAllocator trackingAllocator, long capacity, long memoryAddress) { + // Retain only after wrapForeignAllocation succeeds. On the allocator-limit OOM path, + // wrapForeignAllocation throws before the ForeignAllocation is associated, so release0() + // is not called; retaining first would leave the count elevated with no matching release0(). + ArrowBuf buf = + trackingAllocator.wrapForeignAllocation( + new ForeignAllocation(capacity, memoryAddress) { + @Override + protected void release0() { + ReferenceCountedArrowArray.this.release(); + } + }); retain(); - return trackingAllocator.wrapForeignAllocation( - new ForeignAllocation(capacity, memoryAddress) { - @Override - protected void release0() { - ReferenceCountedArrowArray.this.release(); - } - }); + return buf; } } diff --git a/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java b/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java index f712b400bf..46c93f5541 100644 --- a/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java +++ b/c/src/main/java/org/apache/arrow/c/jni/JniLoader.java @@ -75,8 +75,23 @@ private synchronized void loadRemaining() { } private void load(String name) { - final String libraryToLoad = - name + "/" + getNormalizedArch() + "/" + System.mapLibraryName(name); + String libraryName = System.mapLibraryName(name); + + // If 'arrow.cdata.library.path' is defined, try to load the native library from there + String libraryPath = System.getProperty("arrow.cdata.library.path"); + if (libraryPath != null) { + try { + File libraryFile = new File(libraryPath, libraryName); + if (libraryFile.isFile()) { + System.load(libraryFile.getAbsolutePath()); + return; + } + } catch (UnsatisfiedLinkError e) { + // Ignore this error and fall back to extracting from the JAR file + } + } + + final String libraryToLoad = name + "/" + getNormalizedArch() + "/" + libraryName; try { File temp = File.createTempFile("jnilib-", ".tmp", new File(System.getProperty("java.io.tmpdir"))); diff --git a/c/src/test/java/org/apache/arrow/c/ExceptionTest.java b/c/src/test/java/org/apache/arrow/c/ExceptionTest.java new file mode 100644 index 0000000000..5bc96a8f99 --- /dev/null +++ b/c/src/test/java/org/apache/arrow/c/ExceptionTest.java @@ -0,0 +1,150 @@ +/* + * 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.arrow.c; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowableOfType; + +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorLoader; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.dictionary.DictionaryProvider; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Test; + +// Regression test for https://github.com/apache/arrow-java/issues/759 +final class ExceptionTest { + @Test + public void testException() throws IOException { + final Schema schema = + new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true)))); + final List batches = new ArrayList<>(); + + try (BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + + final String exceptionMessage = "This is a message for testing exception."; + + RuntimeException exToThrow = new RuntimeException(exceptionMessage); + batches.add(exToThrow); + + StringWriter sw = new StringWriter(); + PrintWriter pw = new PrintWriter(sw); + exToThrow.printStackTrace(pw); + final String expectExceptionMessage = sw.toString(); + + ArrowReader source = new ExceptionMemoryArrowReader(allocator, schema, batches); + + try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); + final VectorSchemaRoot importRoot = VectorSchemaRoot.create(schema, allocator)) { + final VectorLoader loader = new VectorLoader(importRoot); + Data.exportArrayStream(allocator, source, stream); + + try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) { + IOException jniException = catchThrowableOfType(IOException.class, reader::loadNextBatch); + final String jniMessage = jniException.getMessage(); + assertThat(jniMessage.endsWith(expectExceptionMessage + "}")); + } + } + } + } + + static class ExceptionMemoryArrowReader extends ArrowReader { + private final Schema schema; + private final List batches; // set ArrowRecordBatch or Exception + private final DictionaryProvider provider; + private int nextBatch; + + ExceptionMemoryArrowReader(BufferAllocator allocator, Schema schema, List batches) { + super(allocator); + this.schema = schema; + this.batches = batches; + this.provider = new CDataDictionaryProvider(); + this.nextBatch = 0; + } + + @Override + public Dictionary lookup(long id) { + return provider.lookup(id); + } + + @Override + public Set getDictionaryIds() { + return provider.getDictionaryIds(); + } + + @Override + public Map getDictionaryVectors() { + return getDictionaryIds().stream() + .collect(Collectors.toMap(Function.identity(), this::lookup)); + } + + @Override + public boolean loadNextBatch() throws IOException { + if (nextBatch < batches.size()) { + Object object = batches.get(nextBatch++); + if (object instanceof RuntimeException) { + throw (RuntimeException) object; + } + VectorLoader loader = new VectorLoader(getVectorSchemaRoot()); + loader.load((ArrowRecordBatch) object); + return true; + } + return false; + } + + @Override + public long bytesRead() { + return 0; + } + + @Override + protected void closeReadSource() throws IOException { + try { + for (Object object : batches) { + if (object instanceof ArrowRecordBatch) { + ArrowRecordBatch batch = (ArrowRecordBatch) object; + batch.close(); + } + } + } catch (Exception e) { + throw new IOException(e); + } + } + + @Override + protected Schema readSchema() { + return schema; + } + } +} diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java new file mode 100644 index 0000000000..7c099f2ef0 --- /dev/null +++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java @@ -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.arrow.c; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.OutOfMemoryException; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Regression test: a mid-import {@link OutOfMemoryException} must not leak the imported array. + * + *

A "producer" allocator owns the exported batch; if the C Data release callback fires, the + * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the + * import. The test asserts the producer drains, confirming the release callback fired despite the + * failure. + */ +final class ImportOutOfMemoryTest { + private static final int ROWS = 1024; + private static final int VALUE_BYTES = 256; + private static final int COLUMNS = 4; + // Far smaller than the exported batch, so the import OOMs part-way through the buffers. + private static final long TINY_LIMIT = 16 * 1024; + + private RootAllocator root; + + @BeforeEach + public void setUp() { + root = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void tearDown() { + root.close(); + } + + @Test + public void importOomDoesNotLeakExportedArray() { + // "producer" owns only the exported batch buffers; the C Data struct containers live on a + // separate allocator (they are consumed/closed by import, which would otherwise muddy the + // producer's balance). So producer draining to zero is an exact signal that the array's release + // callback fired. + try (BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); + BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE)) { + try (ArrowArray array = ArrowArray.allocateNew(structs); + ArrowSchema schema = ArrowSchema.allocateNew(structs)) { + exportBatch(producer, array, schema); + assertTrue( + producer.getAllocatedMemory() > 0, "producer holds the exported batch before import"); + + // A consumer allocator far too small to hold the batch: the import throws part-way through. + try (BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT); + CDataDictionaryProvider provider = new CDataDictionaryProvider()) { + Schema importSchema = Data.importSchema(consumer, schema, provider); + try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) { + Exception thrown = + assertThrows( + Exception.class, + () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider)); + assertTrue( + hasOutOfMemoryCause(thrown), + "mid-import failure must be an allocator OOM: " + thrown); + } + } + + // The array's release callback must have fired despite the mid-import OOM, freeing the + // whole exported batch. On the unfixed retain-before-wrap code the batch is stranded. + assertEquals( + 0L, + producer.getAllocatedMemory(), + "import OOM leaked the exported batch (producer not drained)"); + } + } + } + + /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */ + private static boolean hasOutOfMemoryCause(Throwable t) { + for (Throwable cause = t; cause != null; cause = cause.getCause()) { + if (cause instanceof OutOfMemoryException) { + return true; + } + } + return false; + } + + /** + * Builds a wide multi-column VarChar batch on {@code alloc} and exports it into the C structs. + */ + private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema schema) { + byte[] value = new byte[VALUE_BYTES]; + for (int i = 0; i < value.length; i++) { + value[i] = (byte) 'x'; + } + List vectors = new ArrayList<>(COLUMNS); + for (int c = 0; c < COLUMNS; c++) { + VarCharVector vector = new VarCharVector("col" + c, alloc); + vector.allocateNew((long) ROWS * VALUE_BYTES, ROWS); + for (int r = 0; r < ROWS; r++) { + vector.setSafe(r, value); + } + vector.setValueCount(ROWS); + vectors.add(vector); + } + try (VectorSchemaRoot source = new VectorSchemaRoot(vectors)) { + long total = 0; + for (FieldVector vector : source.getFieldVectors()) { + total += vector.getBufferSize(); + } + assertTrue(total > TINY_LIMIT, "test setup: batch must exceed the consumer limit"); + Data.exportVectorSchemaRoot(alloc, source, null, array, schema); + } + } +} diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java index 6d68449c0b..f6ff88571e 100644 --- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java +++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java @@ -17,9 +17,7 @@ package org.apache.arrow.c; import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.*; import java.nio.ByteBuffer; import java.nio.ByteOrder; @@ -37,7 +35,6 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BigIntVector; @@ -46,7 +43,6 @@ import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; -import org.apache.arrow.vector.ExtensionTypeVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.Float2Vector; @@ -76,6 +72,7 @@ import org.apache.arrow.vector.UInt2Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; @@ -94,6 +91,7 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.UnionMapWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.IntervalDayHolder; import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; @@ -102,7 +100,6 @@ import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; -import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; @@ -812,9 +809,8 @@ public void testEmptyRunEndEncodedVector() { @Test public void testExtensionTypeVector() { - ExtensionTypeRegistry.register(new UuidType()); final Schema schema = - new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); + new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE))); try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { // Fill with data UUID u1 = UUID.randomUUID(); @@ -832,13 +828,12 @@ public void testExtensionTypeVector() { assertEquals(root.getSchema(), importedRoot.getSchema()); final Field field = importedRoot.getSchema().getFields().get(0); - final UuidType expectedType = new UuidType(); assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), - expectedType.extensionName()); + UuidType.INSTANCE.extensionName()); assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_METADATA), - expectedType.serialize()); + UuidType.INSTANCE.serialize()); final UuidVector deserialized = (UuidVector) importedRoot.getFieldVectors().get(0); assertEquals(vector.getValueCount(), deserialized.getValueCount()); @@ -958,6 +953,50 @@ public void testVectorSchemaRootWithDuplicatedFieldNames() { @Test public void testSchema() { + Schema schema = createSchema(); + // Consumer allocates empty ArrowSchema + try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) { + // Producer fills the schema with data + exportSchema(schema, consumerArrowSchema); + + // Consumer imports schema + Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null); + assertEquals(schema.toJson(), importedSchema.toJson()); + } + } + + @Test + public void testSchemaStructReuse() { + Schema schema = createSchema(); + // Consumer allocates empty ArrowSchema + try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) { + // Producer fills the schema with data + exportSchema(schema, consumerArrowSchema); + + // Consumer imports schema + Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null, false); + assertEquals(schema.toJson(), importedSchema.toJson()); + + // Imported struct should be released but not closed + assertEquals(0, consumerArrowSchema.snapshot().release); + assertNotEquals(0, consumerArrowSchema.memoryAddress()); + + // Export and import again + exportSchema(schema, consumerArrowSchema); + importedSchema = Data.importSchema(allocator, consumerArrowSchema, null, false); + assertEquals(schema.toJson(), importedSchema.toJson()); + assertEquals(0, consumerArrowSchema.snapshot().release); + assertNotEquals(0, consumerArrowSchema.memoryAddress()); + } + } + + private void exportSchema(Schema schema, ArrowSchema targetArrowSchema) { + try (ArrowSchema arrowSchema = ArrowSchema.wrap(targetArrowSchema.memoryAddress())) { + Data.exportSchema(allocator, schema, null, arrowSchema); + } + } + + private static Schema createSchema() { Field decimalField = new Field("inner1", FieldType.nullable(new ArrowType.Decimal(19, 4, 128)), null); Field strField = new Field("inner2", FieldType.nullable(new ArrowType.Utf8()), null); @@ -968,16 +1007,7 @@ public void testSchema() { Arrays.asList(decimalField, strField)); Field intField = new Field("col2", FieldType.nullable(new ArrowType.Int(32, true)), null); Schema schema = new Schema(Arrays.asList(itemField, intField)); - // Consumer allocates empty ArrowSchema - try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator)) { - // Producer fills the schema with data - try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress())) { - Data.exportSchema(allocator, schema, null, arrowSchema); - } - // Consumer imports schema - Schema importedSchema = Data.importSchema(allocator, consumerArrowSchema, null); - assertEquals(schema.toJson(), importedSchema.toJson()); - } + return schema; } @Test @@ -1002,12 +1032,8 @@ public void testImportReleasedArray() { try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { // Producer creates structures from existing memory pointers - try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); - ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { - // Producer exports vector into the C Data Interface structures - try (final NullVector vector = new NullVector()) { - Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); - } + try (final NullVector vector = new NullVector()) { + exportFieldVector(vector, consumerArrowSchema, consumerArrowArray); } // Release array structure @@ -1025,6 +1051,45 @@ public void testImportReleasedArray() { } } + @Test + public void testArrayStructReuse() { + // Consumer allocates empty structures + try (ArrowSchema consumerArrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray consumerArrowArray = ArrowArray.allocateNew(allocator)) { + // Producer creates structures from existing memory pointers + try (final NullVector vector = new NullVector()) { + exportFieldVector(vector, consumerArrowSchema, consumerArrowArray); + } + Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null, false); + + // Imported structs should be released but not closed + assertEquals(0, consumerArrowSchema.snapshot().release); + assertNotEquals(0, consumerArrowSchema.memoryAddress()); + assertEquals(0, consumerArrowArray.snapshot().release); + assertNotEquals(0, consumerArrowArray.memoryAddress()); + + try (final NullVector vector = new NullVector()) { + exportFieldVector(vector, consumerArrowSchema, consumerArrowArray); + } + Data.importVector(allocator, consumerArrowArray, consumerArrowSchema, null, false); + + // Imported structs should be released but not closed + assertEquals(0, consumerArrowSchema.snapshot().release); + assertNotEquals(0, consumerArrowSchema.memoryAddress()); + assertEquals(0, consumerArrowArray.snapshot().release); + assertNotEquals(0, consumerArrowArray.memoryAddress()); + } + } + + private void exportFieldVector( + FieldVector vector, ArrowSchema consumerArrowSchema, ArrowArray consumerArrowArray) { + try (ArrowSchema arrowSchema = ArrowSchema.wrap(consumerArrowSchema.memoryAddress()); + ArrowArray arrowArray = ArrowArray.wrap(consumerArrowArray.memoryAddress())) { + // Producer exports vector into the C Data Interface structures + Data.exportVector(allocator, vector, null, arrowArray, arrowSchema); + } + } + private VectorSchemaRoot createTestVSR() { BitVector bitVector = new BitVector("boolean", allocator); @@ -1047,72 +1112,4 @@ private VectorSchemaRoot createTestVSR() { return new VectorSchemaRoot(fields, vectors); } - - static class UuidType extends ExtensionType { - - @Override - public ArrowType storageType() { - return new ArrowType.FixedSizeBinary(16); - } - - @Override - public String extensionName() { - return "uuid"; - } - - @Override - public boolean extensionEquals(ExtensionType other) { - return other instanceof UuidType; - } - - @Override - public ArrowType deserialize(ArrowType storageType, String serializedData) { - if (!storageType.equals(storageType())) { - throw new UnsupportedOperationException( - "Cannot construct UuidType from underlying type " + storageType); - } - return new UuidType(); - } - - @Override - public String serialize() { - return ""; - } - - @Override - public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { - return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16)); - } - } - - static class UuidVector extends ExtensionTypeVector { - - public UuidVector( - String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) { - super(name, allocator, underlyingVector); - } - - @Override - public UUID getObject(int index) { - final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index)); - return new UUID(bb.getLong(), bb.getLong()); - } - - @Override - public int hashCode(int index) { - return hashCode(index, null); - } - - @Override - public int hashCode(int index, ArrowBufHasher hasher) { - return getUnderlyingVector().hashCode(index, hasher); - } - - public void set(int index, UUID uuid) { - ByteBuffer bb = ByteBuffer.allocate(16); - bb.putLong(uuid.getMostSignificantBits()); - bb.putLong(uuid.getLeastSignificantBits()); - getUnderlyingVector().set(index, bb.array()); - } - } } diff --git a/ci/docker/conda-jni.dockerfile b/ci/docker/conda-jni.dockerfile index e14db73688..3f31b74052 100644 --- a/ci/docker/conda-jni.dockerfile +++ b/ci/docker/conda-jni.dockerfile @@ -17,7 +17,7 @@ FROM ghcr.io/mamba-org/micromamba:ubuntu24.04 -ARG jdk=11 +ARG jdk=17 ARG maven=3.9.9 RUN micromamba install -y \ diff --git a/ci/docker/vcpkg-jni.dockerfile b/ci/docker/vcpkg-jni.dockerfile index 55fa35e0d1..f2f5d0d45a 100644 --- a/ci/docker/vcpkg-jni.dockerfile +++ b/ci/docker/vcpkg-jni.dockerfile @@ -18,24 +18,10 @@ ARG base FROM ${base} -# Install the libraries required by Gandiva to run -# Use enable llvm[enable-rtti] in the vcpkg.json to avoid link problems in Gandiva -RUN vcpkg install \ - --clean-after-build \ - --x-install-root=${VCPKG_ROOT}/installed \ - --x-manifest-root=/arrow/ci/vcpkg \ - --x-feature=dev \ - --x-feature=flight \ - --x-feature=gcs \ - --x-feature=json \ - --x-feature=parquet \ - --x-feature=gandiva \ - --x-feature=s3 - # Install Java # We need Java for JNI headers, but we don't invoke Maven in this build. -ARG java=11 -RUN yum install -y java-$java-openjdk-devel && yum clean all +ARG java=17 +RUN dnf install -y java-$java-openjdk-devel && dnf clean all # For ci/scripts/{cpp,java}_*.sh ENV ARROW_HOME=/tmp/local \ diff --git a/ci/scripts/jni_build.sh b/ci/scripts/jni_build.sh index aec6fc325c..c000837987 100755 --- a/ci/scripts/jni_build.sh +++ b/ci/scripts/jni_build.sh @@ -66,7 +66,7 @@ cmake \ -DProtobuf_USE_STATIC_LIBS=ON \ -GNinja \ "${EXTRA_CMAKE_OPTIONS[@]}" -cmake --build "${build_dir}" +cmake --build "${build_dir}" --verbose if [ "${ARROW_JAVA_BUILD_TESTS}" = "ON" ]; then ctest \ --output-on-failure \ diff --git a/ci/scripts/jni_full_build.sh b/ci/scripts/jni_full_build.sh index e9ad0ddbda..5d0aee0555 100755 --- a/ci/scripts/jni_full_build.sh +++ b/ci/scripts/jni_full_build.sh @@ -97,8 +97,10 @@ find ~/.m2/repository/org/apache/arrow \ -exec echo "{}" ";" \ -exec cp "{}" "${dist_dir}" ";" -for artifact in "${dist_dir}"/*; do +pushd "${dist_dir}" +for artifact in *; do sha256sum "${artifact}" >"${artifact}.sha256" sha512sum "${artifact}" >"${artifact}.sha512" done +popd github_actions_group_end diff --git a/ci/scripts/jni_macos_build.sh b/ci/scripts/jni_macos_build.sh index f7543b6f7a..65ab450666 100755 --- a/ci/scripts/jni_macos_build.sh +++ b/ci/scripts/jni_macos_build.sh @@ -59,73 +59,27 @@ fi github_actions_group_begin "Building Arrow C++ libraries" install_dir="${build_dir}/cpp-install" -: "${ARROW_ACERO:=ON}" -export ARROW_ACERO -: "${ARROW_BUILD_TESTS:=OFF}" -export ARROW_BUILD_TESTS -: "${ARROW_DATASET:=ON}" -export ARROW_DATASET -: "${ARROW_GANDIVA:=ON}" -export ARROW_GANDIVA -: "${ARROW_ORC:=ON}" -export ARROW_ORC -: "${ARROW_PARQUET:=ON}" -: "${ARROW_S3:=ON}" -: "${CMAKE_BUILD_TYPE:=Release}" -: "${CMAKE_UNITY_BUILD:=ON}" -export ARROW_TEST_DATA="${arrow_dir}/testing/data" -export PARQUET_TEST_DATA="${arrow_dir}/cpp/submodules/parquet-testing/data" +export ARROW_BUILD_TESTS=OFF + +export ARROW_DATASET=ON +export ARROW_GANDIVA=ON +export ARROW_ORC=ON +export ARROW_PARQUET=ON + export AWS_EC2_METADATA_DISABLED=TRUE cmake \ -S "${arrow_dir}/cpp" \ -B "${build_dir}/cpp" \ - -DARROW_ACERO="${ARROW_ACERO}" \ - -DARROW_BUILD_SHARED=OFF \ - -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \ - -DARROW_CSV="${ARROW_DATASET}" \ - -DARROW_DATASET="${ARROW_DATASET}" \ - -DARROW_SUBSTRAIT="${ARROW_DATASET}" \ - -DARROW_DEPENDENCY_USE_SHARED=OFF \ - -DARROW_GANDIVA="${ARROW_GANDIVA}" \ - -DARROW_GANDIVA_STATIC_LIBSTDCPP=ON \ - -DARROW_JSON="${ARROW_DATASET}" \ - -DARROW_ORC="${ARROW_ORC}" \ - -DARROW_PARQUET="${ARROW_PARQUET}" \ - -DARROW_S3="${ARROW_S3}" \ - -DARROW_USE_CCACHE="${ARROW_USE_CCACHE}" \ - -DAWSSDK_SOURCE=BUNDLED \ - -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \ - -DCMAKE_INSTALL_PREFIX="${install_dir}" \ - -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD}" \ - -DGTest_SOURCE=BUNDLED \ - -DPARQUET_BUILD_EXAMPLES=OFF \ - -DPARQUET_BUILD_EXECUTABLES=OFF \ - -DPARQUET_REQUIRE_ENCRYPTION=OFF \ - -Dre2_SOURCE=BUNDLED \ - -GNinja + --preset=ninja-release-jni-macos \ + -DCMAKE_INSTALL_PREFIX="${install_dir}" cmake --build "${build_dir}/cpp" --target install github_actions_group_end -if [ "${ARROW_RUN_TESTS:-}" == "ON" ]; then - github_actions_group_begin "Running Arrow C++ libraries tests" - # MinIO is required - exclude_tests="arrow-s3fs-test" - # unstable - exclude_tests="${exclude_tests}|arrow-acero-asof-join-node-test" - exclude_tests="${exclude_tests}|arrow-acero-hash-join-node-test" - ctest \ - --exclude-regex "${exclude_tests}" \ - --label-regex unittest \ - --output-on-failure \ - --parallel "$(sysctl -n hw.ncpu)" \ - --test-dir "${build_dir}/cpp" \ - --timeout 300 - github_actions_group_end -fi - -export JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/protobuf_ep-install" +JAVA_JNI_CMAKE_ARGS="-DProtobuf_ROOT=${build_dir}/cpp/_deps/protobuf-build" +JAVA_JNI_CMAKE_ARGS+=" -DProtobuf_SRC_ROOT_FOLDER=${build_dir}/cpp/_deps/protobuf-src" +export JAVA_JNI_CMAKE_ARGS "${source_dir}/ci/scripts/jni_build.sh" \ "${source_dir}" \ "${install_dir}" \ @@ -142,6 +96,7 @@ github_actions_group_begin "Checking shared dependencies for libraries" pushd "${dist_dir}" archery linking check-dependencies \ --allow CoreFoundation \ + --allow Network \ --allow Security \ --allow libSystem \ --allow libarrow_cdata_jni \ diff --git a/ci/scripts/jni_manylinux_build.sh b/ci/scripts/jni_manylinux_build.sh index a34ec0f420..3577c37ab3 100755 --- a/ci/scripts/jni_manylinux_build.sh +++ b/ci/scripts/jni_manylinux_build.sh @@ -53,33 +53,19 @@ if [ "${ARROW_USE_CCACHE}" == "ON" ]; then fi github_actions_group_begin "Building Arrow C++ libraries" -devtoolset_version="$(rpm -qa "devtoolset-*-gcc" --queryformat '%{VERSION}' | grep -o "^[0-9]*")" -devtoolset_include_cpp="/opt/rh/devtoolset-${devtoolset_version}/root/usr/include/c++/${devtoolset_version}" -: "${ARROW_ACERO:=ON}" -export ARROW_ACERO -: "${ARROW_BUILD_TESTS:=OFF}" -export ARROW_BUILD_TESTS -: "${ARROW_DATASET:=ON}" -export ARROW_DATASET -: "${ARROW_GANDIVA:=ON}" -export ARROW_GANDIVA -: "${ARROW_GCS:=ON}" -: "${ARROW_JEMALLOC:=OFF}" -: "${ARROW_MIMALLOC:=ON}" -: "${ARROW_RPATH_ORIGIN:=ON}" -: "${ARROW_ORC:=ON}" -export ARROW_ORC -: "${ARROW_PARQUET:=ON}" -: "${ARROW_S3:=ON}" -: "${CMAKE_BUILD_TYPE:=release}" -: "${CMAKE_UNITY_BUILD:=ON}" + : "${VCPKG_ROOT:=/opt/vcpkg}" : "${VCPKG_FEATURE_FLAGS:=-manifests}" -: "${VCPKG_TARGET_TRIPLET:=${VCPKG_DEFAULT_TRIPLET:-x64-linux-static-${CMAKE_BUILD_TYPE}}}" -: "${GANDIVA_CXX_FLAGS:=-isystem;${devtoolset_include_cpp};-isystem;${devtoolset_include_cpp}/x86_64-redhat-linux;-lpthread}" +: "${VCPKG_TARGET_TRIPLET:=${VCPKG_DEFAULT_TRIPLET:-x64-linux-static-release}}" +export VCPKG_TARGET_TRIPLET + +export ARROW_BUILD_TESTS=OFF + +export ARROW_DATASET=ON +export ARROW_GANDIVA=ON +export ARROW_ORC=ON +export ARROW_PARQUET=ON -export ARROW_TEST_DATA="${arrow_dir}/testing/data" -export PARQUET_TEST_DATA="${arrow_dir}/cpp/submodules/parquet-testing/data" export AWS_EC2_METADATA_DISABLED=TRUE install_dir="${build_dir}/cpp-install" @@ -87,71 +73,12 @@ install_dir="${build_dir}/cpp-install" cmake \ -S "${arrow_dir}/cpp" \ -B "${build_dir}/cpp" \ - -DARROW_ACERO="${ARROW_ACERO}" \ - -DARROW_BUILD_SHARED=OFF \ - -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \ - -DARROW_CSV="${ARROW_DATASET}" \ - -DARROW_DATASET="${ARROW_DATASET}" \ - -DARROW_SUBSTRAIT="${ARROW_DATASET}" \ - -DARROW_DEPENDENCY_SOURCE="VCPKG" \ - -DARROW_DEPENDENCY_USE_SHARED=OFF \ - -DARROW_GANDIVA_PC_CXX_FLAGS="${GANDIVA_CXX_FLAGS}" \ - -DARROW_GANDIVA="${ARROW_GANDIVA}" \ - -DARROW_GCS="${ARROW_GCS}" \ - -DARROW_JEMALLOC="${ARROW_JEMALLOC}" \ - -DARROW_JSON="${ARROW_DATASET}" \ - -DARROW_MIMALLOC="${ARROW_MIMALLOC}" \ - -DARROW_ORC="${ARROW_ORC}" \ - -DARROW_PARQUET="${ARROW_PARQUET}" \ - -DARROW_RPATH_ORIGIN="${ARROW_RPATH_ORIGIN}" \ - -DARROW_S3="${ARROW_S3}" \ - -DARROW_USE_CCACHE="${ARROW_USE_CCACHE}" \ - -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \ - -DCMAKE_INSTALL_PREFIX="${install_dir}" \ - -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD}" \ - -DGTest_SOURCE=BUNDLED \ - -DORC_SOURCE=BUNDLED \ - -DORC_PROTOBUF_EXECUTABLE="${VCPKG_ROOT}/installed/${VCPKG_TARGET_TRIPLET}/tools/protobuf/protoc" \ - -DPARQUET_BUILD_EXAMPLES=OFF \ - -DPARQUET_BUILD_EXECUTABLES=OFF \ - -DPARQUET_REQUIRE_ENCRYPTION=OFF \ - -DVCPKG_MANIFEST_MODE=OFF \ - -DVCPKG_TARGET_TRIPLET="${VCPKG_TARGET_TRIPLET}" \ - -GNinja + --preset=ninja-release-jni-linux \ + -DCMAKE_INSTALL_PREFIX="${install_dir}" cmake --build "${build_dir}/cpp" cmake --install "${build_dir}/cpp" github_actions_group_end -if [ "${ARROW_RUN_TESTS:-OFF}" = "ON" ]; then - github_actions_group_begin "Running Arrow C++ libraries tests" - # MinIO is required - exclude_tests="arrow-s3fs-test" - case $(arch) in - aarch64) - # GCS testbench is crashed on aarch64: - # ImportError: ../grpc/_cython/cygrpc.cpython-38-aarch64-linux-gnu.so: - # undefined symbol: vtable for std::__cxx11::basic_ostringstream< - # char, std::char_traits, std::allocator > - exclude_tests="${exclude_tests}|arrow-gcsfs-test" - ;; - esac - # unstable - exclude_tests="${exclude_tests}|arrow-acero-asof-join-node-test" - exclude_tests="${exclude_tests}|arrow-acero-hash-join-node-test" - # external dependency - exclude_tests="${exclude_tests}|arrow-gcsfs-test" - # strptime - exclude_tests="${exclude_tests}|arrow-utility-test" - ctest \ - --exclude-regex "${exclude_tests}" \ - --label-regex unittest \ - --output-on-failure \ - --parallel "$(nproc)" \ - --test-dir "${build_dir}/cpp" \ - --timeout 300 - github_actions_group_end -fi - JAVA_JNI_CMAKE_ARGS="-DCMAKE_TOOLCHAIN_FILE=${VCPKG_ROOT}/scripts/buildsystems/vcpkg.cmake" JAVA_JNI_CMAKE_ARGS="${JAVA_JNI_CMAKE_ARGS} -DVCPKG_TARGET_TRIPLET=${VCPKG_TARGET_TRIPLET}" export JAVA_JNI_CMAKE_ARGS diff --git a/ci/scripts/jni_windows_build.sh b/ci/scripts/jni_windows_build.sh index d01ef45f5a..6503ac63e5 100755 --- a/ci/scripts/jni_windows_build.sh +++ b/ci/scripts/jni_windows_build.sh @@ -68,7 +68,7 @@ cmake \ -B "${build_dir}/cpp" \ -DARROW_ACERO="${ARROW_ACERO}" \ -DARROW_BUILD_SHARED=OFF \ - -DARROW_BUILD_TESTS=ON \ + -DARROW_BUILD_TESTS="${ARROW_BUILD_TESTS}" \ -DARROW_CSV="${ARROW_DATASET}" \ -DARROW_DATASET="${ARROW_DATASET}" \ -DARROW_SUBSTRAIT="${ARROW_DATASET}" \ diff --git a/ci/scripts/test.sh b/ci/scripts/test.sh index cacc20034e..8061ee455d 100755 --- a/ci/scripts/test.sh +++ b/ci/scripts/test.sh @@ -34,10 +34,11 @@ fi mvn="mvn -B -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn" # Use `2 * ncores` threads mvn="${mvn} -T 2C" +mvn="${mvn} -Denforcer.skip=true" pushd "${build_dir}" -${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" clean test +${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" test projects=() if [ "${ARROW_JAVA_JNI}" = "ON" ]; then @@ -46,7 +47,7 @@ if [ "${ARROW_JAVA_JNI}" = "ON" ]; then projects+=(gandiva) fi if [ "${#projects[@]}" -gt 0 ]; then - ${mvn} clean test \ + ${mvn} test \ -Parrow-jni \ -pl "$( IFS=, @@ -56,7 +57,7 @@ if [ "${#projects[@]}" -gt 0 ]; then fi if [ "${ARROW_JAVA_CDATA}" = "ON" ]; then - ${mvn} clean test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}" + ${mvn} test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}" fi popd diff --git a/compose.yaml b/compose.yaml index b125c3c983..4fd825e5a5 100644 --- a/compose.yaml +++ b/compose.yaml @@ -40,7 +40,7 @@ services: # docker compose run ubuntu # Parameters: # MAVEN: 3.9.9 - # JDK: 11, 17, 21 + # JDK: 17, 21 image: ${ARCH}/maven:${MAVEN}-eclipse-temurin-${JDK} volumes: - .:/arrow-java:delegated @@ -60,7 +60,7 @@ services: # docker compose run conda-jni-cdata # Parameters: # MAVEN: 3.9.9 - # JDK: 11, 17, 21 + # JDK: 17, 21 image: ${REPO}:${ARCH}-conda-java-${JDK}-maven-${MAVEN}-jni-integration build: context: . @@ -99,7 +99,7 @@ services: cache_from: - ${REPO}:${ARCH}-vcpkg-jni-${VCPKG} args: - base: ${ARROW_REPO}:${ARCH}-python-${PYTHON}-wheel-manylinux-2014-vcpkg-${VCPKG} + base: ${ARROW_REPO}:${ARCH}-cpp-jni-${VCPKG} volumes: - .:/arrow-java:delegated - ${ARROW_REPO_ROOT}:/arrow:delegated @@ -109,5 +109,6 @@ services: ARROW_JAVA_CDATA: "ON" CCACHE_DIR: "/ccache" command: - ["git config --global --add safe.directory /arrow-java && \ + ["/bin/bash", "-c", + "git config --global --add safe.directory /arrow-java && /arrow-java/ci/scripts/jni_manylinux_build.sh /arrow-java /arrow /build/java /arrow-java/jni"] diff --git a/compression/pom.xml b/compression/pom.xml index 6f60eb7d0a..aa7dee6f89 100644 --- a/compression/pom.xml +++ b/compression/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-compression Arrow Compression @@ -50,12 +50,12 @@ under the License. org.apache.commons commons-compress - 1.27.1 + 1.28.0 com.github.luben zstd-jni - 1.5.7-2 + 1.5.7-11 diff --git a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java index 91cefc2a9e..f268e815fe 100644 --- a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java +++ b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java @@ -80,6 +80,13 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu } byte[] outBytes = out.toByteArray(); + if (outBytes.length != decompressedLength) { + throw new RuntimeException( + "Expected != actual decompressed length: " + + decompressedLength + + " != " + + outBytes.length); + } ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length); decompressedBuffer.setBytes(/* index= */ 0, outBytes); decompressedBuffer.writerIndex(decompressedLength); diff --git a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java index 290723608d..ed46fe81b4 100644 --- a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java +++ b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java @@ -44,7 +44,7 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu long bytesWritten = Zstd.compressUnsafe( compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, - dstSize, + maxSize, /*src*/ uncompressedBuffer.memoryAddress(), /* srcSize= */ uncompressedBuffer.writerIndex(), /* level= */ this.compressionLevel); diff --git a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java index b8fb4e28b9..d2d2921649 100644 --- a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java +++ b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java @@ -20,6 +20,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayOutputStream; @@ -59,6 +60,7 @@ import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; @@ -231,6 +233,26 @@ void testEmptyBuffer(int vectorLength, CompressionCodec codec) throws Exception AutoCloseables.close(decompressedBuffers); } + @Test + void testLz4DecompressRejectsWrongLength() { + byte[] data = new byte[512]; // all zeros, highly compressible + ArrowBuf orig = allocator.buffer(data.length); + orig.setBytes(0, data); + orig.writerIndex(data.length); + + CompressionCodec codec = new Lz4CompressionCodec(); + ArrowBuf compressed = codec.compress(allocator, orig); + + // tamper with the 8-byte uncompressed-length prefix so it no longer matches + // the real decompressed size + compressed.setLong(0, 1_000_000L); + + RuntimeException e = + assertThrows(RuntimeException.class, () -> codec.decompress(allocator, compressed)); + assertTrue(e.getMessage().contains("decompressed length")); + compressed.close(); + } + private static Stream codecTypes() { return Arrays.stream(CompressionUtil.CodecType.values()); } diff --git a/dataset/pom.xml b/dataset/pom.xml index 7c2da2ac21..5acc837860 100644 --- a/dataset/pom.xml +++ b/dataset/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-dataset @@ -32,8 +32,8 @@ under the License. ../../../cpp/release-build/ - 1.15.1 - 1.12.0 + 1.17.1 + 1.12.1 @@ -130,7 +130,7 @@ under the License. org.apache.orc orc-core - 2.1.1 + 2.3.0 test @@ -156,7 +156,7 @@ under the License. commons-io commons-io - 2.18.0 + 2.22.0 test diff --git a/dataset/src/main/cpp/jni_wrapper.cc b/dataset/src/main/cpp/jni_wrapper.cc index 49cc85251c..e8087648eb 100644 --- a/dataset/src/main/cpp/jni_wrapper.cc +++ b/dataset/src/main/cpp/jni_wrapper.cc @@ -23,6 +23,7 @@ #include "arrow/array/concatenate.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "arrow/compute/initialize.h" #include "arrow/dataset/api.h" #include "arrow/dataset/file_base.h" #ifdef ARROW_CSV @@ -807,6 +808,13 @@ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_ensureS3Fina JNI_METHOD_END() } +JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_initialize( + JNIEnv* env, jobject) { + JNI_METHOD_START + JniAssertOkOrThrow(arrow::compute::Initialize()); + JNI_METHOD_END() +} + /* * Class: org_apache_arrow_dataset_file_JniWrapper * Method: makeFileSystemDatasetFactory diff --git a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java index 631b8b1bbe..5fb4816488 100644 --- a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java +++ b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java @@ -56,6 +56,7 @@ public void ensureLoaded() { } loadRemaining(); ensureS3FinalizedOnShutdown(); + JniWrapper.get().initialize(); } private synchronized void loadRemaining() { diff --git a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java index 6637c113d9..cfef098ec4 100644 --- a/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java +++ b/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java @@ -124,4 +124,7 @@ public native long createScanner( * uninitialized, then this is a noop. */ public native void ensureS3Finalized(); + + /** Initialize Arrow Compute. */ + public native void initialize(); } diff --git a/dev/release/bump_version.sh b/dev/release/bump_version.sh index 458e930f98..68cafb99bd 100755 --- a/dev/release/bump_version.sh +++ b/dev/release/bump_version.sh @@ -37,6 +37,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then exit 1 fi . "${SOURCE_DIR}/.env" +export GH_TOKEN cd "${SOURCE_TOP_DIR}" diff --git a/dev/release/release.sh b/dev/release/release.sh index 70e1f96454..d1db7ad05a 100755 --- a/dev/release/release.sh +++ b/dev/release/release.sh @@ -19,6 +19,8 @@ set -eu +SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + if [ "$#" -ne 2 ]; then echo "Usage: $0 " echo " e.g.: $0 19.0.1 1" @@ -34,6 +36,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then exit 1 fi . "${SOURCE_DIR}/.env" +export GH_TOKEN git_origin_url="$(git remote get-url origin)" repository="${git_origin_url#*github.com?}" diff --git a/dev/release/release_rc.sh b/dev/release/release_rc.sh index ff77718b8d..0920edbe35 100755 --- a/dev/release/release_rc.sh +++ b/dev/release/release_rc.sh @@ -42,6 +42,7 @@ if [ ! -f "${SOURCE_DIR}/.env" ]; then exit 1 fi . "${SOURCE_DIR}/.env" +export GH_TOKEN cd "${SOURCE_TOP_DIR}" diff --git a/docs/source/cdata.rst b/docs/source/cdata.rst index 9643d88df3..7b2924d259 100644 --- a/docs/source/cdata.rst +++ b/docs/source/cdata.rst @@ -101,8 +101,8 @@ without writing JNI bindings ourselves. 1.0-SNAPSHOT - 8 - 8 + 17 + 17 9.0.0 @@ -237,8 +237,8 @@ For this example, we will build a JAR with all dependencies bundled. cpptojava 1.0-SNAPSHOT - 8 - 8 + 17 + 17 9.0.0 diff --git a/docs/source/developers/building.rst b/docs/source/developers/building.rst index f9ef7daea8..b682957714 100644 --- a/docs/source/developers/building.rst +++ b/docs/source/developers/building.rst @@ -32,7 +32,7 @@ Arrow Java uses the `Maven `_ build system. Building requires: -* JDK 11+ +* JDK 17+ * Maven 3+ .. note:: @@ -345,7 +345,7 @@ configuration file usually located under ``${HOME}/.m2`` with the following snip jdk - 21 + 21 temurin @@ -383,11 +383,11 @@ Arrow repository, and update the following settings: right click the directory, and select Mark Directory as > Generated Sources Root. There is no need to mark other generated sources directories, as only the ``vector`` module generates sources. -* For JDK 11, due to an `IntelliJ bug - `__, you must go into +* Due to an `IntelliJ bug + `__, you may need to go into Settings > Build, Execution, Deployment > Compiler > Java Compiler and disable "Use '--release' option for cross-compilation (Java 9 and later)". Otherwise - you will get an error like "package sun.misc does not exist". + you may get an error like "package sun.misc does not exist". * You may want to disable error-prone entirely if it gives spurious warnings (disable both error-prone profiles in the Maven tool window and "Reload All Maven Projects"). @@ -397,7 +397,7 @@ Arrow repository, and update the following settings: * To enable debugging JNI-based modules like ``dataset``, activate specific profiles in the Maven tab under "Profiles". Ensure the profiles ``arrow-c-data``, ``arrow-jni``, ``generate-libs-cdata-all-os``, - ``generate-libs-jni-macos-linux``, and ``jdk11+`` are enabled, so that the + ``generate-libs-jni-macos-linux``, and ``jdk17+`` are enabled, so that the IDE can build them and enable debugging. You may not need to update all of these settings if you build/test with the @@ -478,8 +478,8 @@ Installing Manually .. code-block:: xml - 8 - 8 + 17 + 17 9.0.0.dev501 diff --git a/docs/source/flight.rst b/docs/source/flight.rst index fabced8094..fd0fdf07bc 100644 --- a/docs/source/flight.rst +++ b/docs/source/flight.rst @@ -232,8 +232,8 @@ Servers can add other gRPC services. For example, to add the `Health Check servi See the :external+arrow:ref:`best practices for C++ `. -.. _`FlightClient`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightClient.html -.. _`FlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightProducer.html -.. _`FlightServer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/FlightServer.html -.. _`NoOpFlightProducer`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/NoOpFlightProducer.html -.. _`Location`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/Location.html +.. _`FlightClient`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightClient.html +.. _`FlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightProducer.html +.. _`FlightServer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/FlightServer.html +.. _`NoOpFlightProducer`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/NoOpFlightProducer.html +.. _`Location`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/Location.html diff --git a/docs/source/flight_sql.rst b/docs/source/flight_sql.rst index 169a0e24bf..09ce1dda0d 100644 --- a/docs/source/flight_sql.rst +++ b/docs/source/flight_sql.rst @@ -29,4 +29,4 @@ over the network. For usage information, see the `API documentation`_. -.. _API documentation: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/sql/package-summary.html +.. _API documentation: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.sql/org/apache/arrow/flight/sql/package-summary.html diff --git a/docs/source/flight_sql_jdbc_driver.rst b/docs/source/flight_sql_jdbc_driver.rst index 1806930943..6d40434a22 100644 --- a/docs/source/flight_sql_jdbc_driver.rst +++ b/docs/source/flight_sql_jdbc_driver.rst @@ -27,7 +27,7 @@ Flight SQL. Installation and Requirements ============================= -The driver is compatible with JDK 11+. Note that the following JVM +The driver is compatible with JDK 17+. Note that the following JVM parameter is required: .. code-block:: shell @@ -173,3 +173,126 @@ DriverManager#getConnection() `_, the username and password supplied on the URI supercede the username and password arguments to the function call. + +OAuth 2.0 Authentication +======================== + +The driver supports OAuth 2.0 authentication for obtaining access tokens +from an authorization server. Two OAuth flows are currently supported: + +* **Client Credentials** - For service-to-service authentication where no + user interaction is required. The application authenticates using its own + credentials (client ID and client secret). + +* **Token Exchange** (RFC 8693) - For exchanging one token for another, + commonly used for federated authentication, delegation, or impersonation + scenarios. + +OAuth Connection Properties +--------------------------- + +The following properties configure OAuth authentication. These properties +should be provided via the ``Properties`` object when connecting, as they +may contain special characters that are difficult to encode in a URI. + +**Common OAuth Properties** + +.. list-table:: + :header-rows: 1 + + * - Parameter + - Type + - Required + - Default + - Description + + * - oauth.flow + - String + - Yes (to enable OAuth) + - null + - The OAuth grant type. Supported values: ``client_credentials``, + ``token_exchange`` + + * - oauth.tokenUri + - String + - Yes + - null + - The OAuth 2.0 token endpoint URL (e.g., + ``https://auth.example.com/oauth/token``) + + * - oauth.clientId + - String + - Conditional + - null + - The OAuth 2.0 client ID. Required for ``client_credentials`` flow, + optional for ``token_exchange`` + + * - oauth.clientSecret + - String + - Conditional + - null + - The OAuth 2.0 client secret. Required for ``client_credentials`` flow, + optional for ``token_exchange`` + + * - oauth.scope + - String + - No + - null + - Space-separated list of OAuth scopes to request + + * - oauth.resource + - String + - No + - null + - The resource indicator for the token request (RFC 8707) + +**Token Exchange Properties** + +These properties are specific to the ``token_exchange`` flow: + +.. list-table:: + :header-rows: 1 + + * - Parameter + - Type + - Required + - Default + - Description + + * - oauth.exchange.subjectToken + - String + - Yes + - null + - The subject token to exchange (e.g., a JWT from an identity provider) + + * - oauth.exchange.subjectTokenType + - String + - Yes + - null + - The token type URI of the subject token. Common values: + ``urn:ietf:params:oauth:token-type:access_token``, + ``urn:ietf:params:oauth:token-type:jwt`` + + * - oauth.exchange.actorToken + - String + - No + - null + - The actor token for delegation/impersonation scenarios + + * - oauth.exchange.actorTokenType + - String + - No + - null + - The token type URI of the actor token + + * - oauth.exchange.aud + - String + - No + - null + - The target audience for the exchanged token + + * - oauth.exchange.requestedTokenType + - String + - No + - null + - The desired token type for the exchanged token diff --git a/docs/source/install.rst b/docs/source/install.rst index b2b1c7163f..e0b34515ef 100644 --- a/docs/source/install.rst +++ b/docs/source/install.rst @@ -27,8 +27,8 @@ Java modules are regularly built and tested on macOS and Linux distributions. Java Compatibility ================== -Java modules are compatible with JDK 11 and above. Currently, JDK versions -11, 17, 21, and latest are tested in CI. +Java modules are compatible with JDK 17 and above. Currently, JDK versions +17, 21, and latest are tested in CI. Note that some JDK internals must be exposed by adding ``--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED`` to the ``java`` command: diff --git a/docs/source/jdbc.rst b/docs/source/jdbc.rst index c0477cb06d..e054127faa 100644 --- a/docs/source/jdbc.rst +++ b/docs/source/jdbc.rst @@ -95,7 +95,7 @@ Type Mapping The JDBC to Arrow type mapping can be obtained at runtime from `JdbcToArrowUtils.getArrowTypeFromJdbcType`_. -.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar- +.. _JdbcToArrowUtils.getArrowTypeFromJdbcType: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.html#getArrowTypeFromJdbcType-org.apache.arrow.adapter.jdbc.JdbcFieldInfo-java.util.Calendar- +--------------------+--------------------+-------+ | JDBC Type | Arrow Type | Notes | @@ -171,8 +171,8 @@ The JDBC to Arrow type mapping can be obtained at runtime from timezone of the calendar, else it will be a timestamp without timezone. -.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map- -.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/docs/java/reference/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map- +.. _setArraySubTypeByColumnIndexMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnIndexMap-java.util.Map- +.. _setArraySubTypeByColumnNameMap: https://arrow.apache.org/java/current/reference/org.apache.arrow.adapter.jdbc/org/apache/arrow/adapter/jdbc/JdbcToArrowConfigBuilder.html#setArraySubTypeByColumnNameMap-java.util.Map- .. _ARROW-17006: https://issues.apache.org/jira/browse/ARROW-17006 VectorSchemaRoot to PreparedStatement Parameter Conversion @@ -213,7 +213,8 @@ Type Mapping ------------ The Arrow to JDBC type mapping can be obtained at runtime via -a method on ColumnBinder. +a method on ColumnBinder. The Flight SQL JDBC driver follows the same +mapping, with additional support for the UUID extension type noted below. +----------------------------+----------------------------+-------+ | Arrow Type | JDBC Type | Notes | @@ -232,6 +233,8 @@ a method on ColumnBinder. +----------------------------+----------------------------+-------+ | FixedSizeBinary | BINARY (setBytes) | | +----------------------------+----------------------------+-------+ +| Uuid (extension) | OTHER (setObject) | \(3) | ++----------------------------+----------------------------+-------+ | Float32 | REAL (setFloat) | | +----------------------------+----------------------------+-------+ | Int8 | TINYINT (setByte) | | @@ -273,6 +276,9 @@ a method on ColumnBinder. JDBC value, because a JDBC Timestamp is in UTC, and we have no timezone information. In this case, the default binder will call `setTimestamp(int, Timestamp) - `_, + `_, which will lead to the driver using the "default timezone" (that of the Java VM). +* \(3) For the Flight SQL JDBC driver, the Arrow UUID extension type + (``arrow.uuid``) maps to JDBC ``OTHER`` and is surfaced as + ``java.util.UUID`` values. diff --git a/docs/source/memory.rst b/docs/source/memory.rst index 58ef382dc9..5b9148223a 100644 --- a/docs/source/memory.rst +++ b/docs/source/memory.rst @@ -333,18 +333,18 @@ How this works: } } -.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html -.. _`ArrowBuf.print()`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity- -.. _`BufferAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferAllocator.html -.. _`BufferLedger`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/BufferLedger.html -.. _`RootAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html -.. _`newChildAllocator`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long- +.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html +.. _`ArrowBuf.print()`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html#print-java.lang.StringBuilder-int-org.apache.arrow.memory.BaseAllocator.Verbosity- +.. _`BufferAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferAllocator.html +.. _`BufferLedger`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/BufferLedger.html +.. _`RootAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html +.. _`newChildAllocator`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html#newChildAllocator-java.lang.String-org.apache.arrow.memory.AllocationListener-long-long- .. _`Netty`: https://netty.io/wiki/ .. _`sun.misc.unsafe`: https://web.archive.org/web/20210929024401/http://www.docjar.com/html/api/sun/misc/Unsafe.java.html -.. _`Direct Memory`: https://docs.oracle.com/en/java/javase/11/docs/api/java.base/java/nio/ByteBuffer.html -.. _`ReferenceManager`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html -.. _`ReferenceManager.release`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#release-- -.. _`ReferenceManager.retain`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ReferenceManager.html#retain-- +.. _`Direct Memory`: https://docs.oracle.com/en/java/javase/17/docs/api/java.base/java/nio/ByteBuffer.html +.. _`ReferenceManager`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html +.. _`ReferenceManager.release`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#release-- +.. _`ReferenceManager.retain`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ReferenceManager.html#retain-- Arrow Memory In-Depth ===================== diff --git a/docs/source/overview.rst b/docs/source/overview.rst index be579c1495..1188054114 100644 --- a/docs/source/overview.rst +++ b/docs/source/overview.rst @@ -45,6 +45,9 @@ but some modules are JNI bindings to the C++ library. * - arrow-vector - An off-heap reference implementation for Arrow columnar data format. - Native + * - arrow-vector-codegen + - Template files for Arrow datatypes suitable for code generation. + - Native * - arrow-tools - Java applications for working with Arrow ValueVectors. - Native diff --git a/docs/source/substrait.rst b/docs/source/substrait.rst index b3678ac815..5ec07f1658 100644 --- a/docs/source/substrait.rst +++ b/docs/source/substrait.rst @@ -19,7 +19,7 @@ Substrait ========= -The ``arrow-dataset`` module can execute Substrait_ plans via the :external+arrow:doc:`Acero ` +The ``arrow-dataset`` module can execute Substrait_ plans via the :external+arrow:doc:`Acero ` query engine. Executing Queries Using Substrait Plans diff --git a/docs/source/table.rst b/docs/source/table.rst index 5aa95e153c..880ef84d29 100644 --- a/docs/source/table.rst +++ b/docs/source/table.rst @@ -364,15 +364,15 @@ If the table contains dictionary-encoded vectors and was constructed with a ``Di Data.exportTable(bufferAllocator, table, outArrowArray); -.. _`ArrowBuf`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/memory/ArrowBuf.html -.. _`Data`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/c/Data.html -.. _`DictionaryProvider`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/dictionary/DictionaryProvider.html -.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html -.. _`FieldReader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/complex/reader/FieldReader.html -.. _`FieldVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/FieldVector.html -.. _`Row`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Row.html -.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html -.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html -.. _`ValueHolder`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/holders/ValueHolder.html -.. _`ValueVector`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ValueVector.html -.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html +.. _`ArrowBuf`: https://arrow.apache.org/java/current/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/ArrowBuf.html +.. _`Data`: https://arrow.apache.org/java/current/reference/org.apache.arrow.c/org/apache/arrow/c/Data.html +.. _`DictionaryProvider`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/dictionary/DictionaryProvider.html +.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html +.. _`FieldReader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/complex/reader/FieldReader.html +.. _`FieldVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/FieldVector.html +.. _`Row`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Row.html +.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html +.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html +.. _`ValueHolder`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/holders/ValueHolder.html +.. _`ValueVector`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ValueVector.html +.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html diff --git a/docs/source/vector_schema_root.rst b/docs/source/vector_schema_root.rst index 3119122d9a..f4a497c4e5 100644 --- a/docs/source/vector_schema_root.rst +++ b/docs/source/vector_schema_root.rst @@ -153,11 +153,11 @@ A `Table`_ is an immutable tabular data structure, very similar to VectorSchemaR See the :doc:`table` documentation for more information. -.. _`ArrowRecordBatch`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html -.. _`Field`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Field.html -.. _`Flight`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/flight/package-summary.html -.. _`Schema`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/types/pojo/Schema.html -.. _`Table`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/table/Table.html -.. _`VectorLoader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorLoader.html -.. _`VectorSchemaRoot`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorSchemaRoot.html -.. _`VectorUnloader`: https://arrow.apache.org/docs/java/reference/org/apache/arrow/vector/VectorUnloader.html +.. _`ArrowRecordBatch`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/ipc/message/ArrowRecordBatch.html +.. _`Field`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Field.html +.. _`Flight`: https://arrow.apache.org/java/current/reference/org.apache.arrow.flight.core/org/apache/arrow/flight/package-summary.html +.. _`Schema`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/types/pojo/Schema.html +.. _`Table`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/table/Table.html +.. _`VectorLoader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorLoader.html +.. _`VectorSchemaRoot`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorSchemaRoot.html +.. _`VectorUnloader`: https://arrow.apache.org/java/current/reference/org.apache.arrow.vector/org/apache/arrow/vector/VectorUnloader.html diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml index a32293a5e8..9ae402cdc4 100644 --- a/flight/flight-core/pom.xml +++ b/flight/flight-core/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-flight - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT flight-core @@ -134,7 +134,7 @@ under the License. com.google.api.grpc proto-google-common-protos - 2.54.1 + 2.72.0 test diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java index 9cefccb3fe..ab4eab3048 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java @@ -287,7 +287,11 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s ArrowBuf body = null; ArrowBuf appMetadata = null; while (stream.available() > 0) { - int tag = readRawVarint32(stream); + final int tagFirstByte = stream.read(); + if (tagFirstByte == -1) { + break; + } + int tag = readRawVarint32(tagFirstByte, stream); switch (tag) { case DESCRIPTOR_TAG: { @@ -366,6 +370,10 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s private static int readRawVarint32(InputStream is) throws IOException { int firstByte = is.read(); + return readRawVarint32(firstByte, is); + } + + private static int readRawVarint32(int firstByte, InputStream is) throws IOException { return CodedInputStream.readRawVarint32(firstByte, is); } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java index f4f6486a3c..0939d232cf 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/CallHeaders.java @@ -26,10 +26,20 @@ public interface CallHeaders { /** Get the value of a metadata key. If multiple values are present, then get the last one. */ byte[] getByte(String key); - /** Get all values present for the given metadata key. */ + /** + * Get all values present for the given metadata key. + * + * @param key the metadata key + * @return an iterable of all values for the key. Returns an empty iterable if no value to return. + */ Iterable getAll(String key); - /** Get all values present for the given metadata key. */ + /** + * Get all values present for the given metadata key. + * + * @param key the metadata key + * @return an iterable of all values for the key. Returns an empty iterable if no value to return. + */ Iterable getAllByte(String key); /** diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java index 47fd6f1366..5ec01b9c83 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java @@ -80,20 +80,18 @@ public ServerSessionMiddleware onCallStarted( String sessionId = null; final Iterable it = incomingHeaders.getAll("cookie"); - if (it != null) { - findIdCookie: - for (final String headerValue : it) { - for (final String cookie : headerValue.split(" ;")) { - final String[] cookiePair = cookie.split("="); - if (cookiePair.length != 2) { - // Soft failure: Ignore invalid cookie list field - break; - } - - if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) { - sessionId = cookiePair[1]; - break findIdCookie; - } + findIdCookie: + for (final String headerValue : it) { + for (final String cookie : headerValue.split(" ;")) { + final String[] cookiePair = cookie.split("="); + if (cookiePair.length != 2) { + // Soft failure: Ignore invalid cookie list field + break; + } + + if (sessionCookieName.equals(cookiePair[0]) && cookiePair[1].length() > 0) { + sessionId = cookiePair[1]; + break findIdCookie; } } } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java index e5eb934001..b33e6b7ecc 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java @@ -100,10 +100,7 @@ public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { @Override public void onHeadersReceived(CallHeaders incomingHeaders) { - final Iterable setCookieHeaders = incomingHeaders.getAll(SET_COOKIE_HEADER); - if (setCookieHeaders != null) { - factory.updateCookies(setCookieHeaders); - } + factory.updateCookies(incomingHeaders.getAll(SET_COOKIE_HEADER)); } @Override diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java index f33e9b2f94..fe81f3fb23 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/CallCredentialAdapter.java @@ -18,6 +18,7 @@ import io.grpc.CallCredentials; import io.grpc.Metadata; +import io.grpc.Status; import java.util.concurrent.Executor; import java.util.function.Consumer; import org.apache.arrow.flight.CallHeaders; @@ -36,9 +37,14 @@ public void applyRequestMetadata( RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) { executor.execute( () -> { - final Metadata headers = new Metadata(); - credentialWriter.accept(new MetadataAdapter(headers)); - metadataApplier.apply(headers); + try { + final Metadata headers = new Metadata(); + credentialWriter.accept(new MetadataAdapter(headers)); + metadataApplier.apply(headers); + } catch (Throwable t) { + metadataApplier.fail( + Status.UNAUTHENTICATED.withCause(t).withDescription(t.getMessage())); + } }); } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index 45c32a86c6..fcba88d212 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -87,13 +87,13 @@ public static void readIntoBuffer( final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath) throws IOException { ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null; + byte[] heapBytes = new byte[size]; if (readableBuffer != null) { - readableBuffer.readBytes(buf.nioBuffer(0, size)); + readableBuffer.readBytes(heapBytes, 0, size); } else { - byte[] heapBytes = new byte[size]; ByteStreams.readFully(stream, heapBytes); - buf.writeBytes(heapBytes); } + buf.writeBytes(heapBytes); buf.writerIndex(size); } } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java index a1de16ede6..64a0769d63 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/MetadataAdapter.java @@ -18,6 +18,7 @@ import io.grpc.Metadata; import java.nio.charset.StandardCharsets; +import java.util.Collections; import java.util.HashSet; import java.util.Set; import java.util.stream.Collectors; @@ -53,13 +54,17 @@ public byte[] getByte(String key) { @Override public Iterable getAll(String key) { - return this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)); + final Iterable all = + this.metadata.getAll(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)); + return all != null ? all : Collections.emptyList(); } @Override public Iterable getAllByte(String key) { if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { - return this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER)); + final Iterable all = + this.metadata.getAll(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER)); + return all != null ? all : Collections.emptyList(); } return StreamSupport.stream(getAll(key).spliterator(), false) .map(String::getBytes) diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java index a54ce69812..8aef9c69a1 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestCallOptions.java @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -110,6 +111,16 @@ public void mixedProperties() { testHeaders(headers); } + @Test + public void getAllReturnsEmptyIterableForMissingKey() { + FlightCallHeaders headers = new FlightCallHeaders(); + + assertNotNull(headers.getAll("missing")); + assertFalse(headers.getAll("missing").iterator().hasNext()); + assertNotNull(headers.getAllByte("missing-bin")); + assertFalse(headers.getAllByte("missing-bin").iterator().hasNext()); + } + private void testHeaders(CallHeaders headers) { try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE); HeaderProducer producer = new HeaderProducer(); diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java index a9a3e355bc..214614defd 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestErrorMetadata.java @@ -20,6 +20,7 @@ import static org.apache.arrow.flight.Location.forGrpcInsecure; import static org.junit.jupiter.api.Assertions.assertArrayEquals; 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.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -119,6 +120,16 @@ public void testFlightMetadata() throws Exception { } } + @Test + public void getAllReturnsEmptyIterableForMissingKey() { + ErrorFlightMetadata metadata = new ErrorFlightMetadata(); + + assertNotNull(metadata.getAll("missing")); + assertFalse(metadata.getAll("missing").iterator().hasNext()); + assertNotNull(metadata.getAllByte("missing-bin")); + assertFalse(metadata.getAllByte("missing-bin").iterator().hasNext()); + } + private static class StatusRuntimeExceptionProducer extends NoOpFlightProducer { private final PerfOuterClass.Perf perf; diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java index 0c63785c88..0f202ba2d9 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java @@ -178,6 +178,12 @@ public static void shutdown() throws Exception { AutoCloseables.close(server); allocator.getChildAllocators().forEach(BufferAllocator::close); + + // gRPC/Netty may still be releasing Arrow buffers asynchronously after server shutdown. + // Poll briefly to allow in-flight buffer releases to complete before closing the allocator. + for (int i = 0; i < 20 && allocator.getAllocatedMemory() > 0; i++) { + Thread.sleep(100); + } AutoCloseables.close(allocator); } } diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java new file mode 100644 index 0000000000..b0f5dcfcfc --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestMetadataAdapter.java @@ -0,0 +1,36 @@ +/* + * 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.arrow.flight.grpc; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import io.grpc.Metadata; +import org.junit.jupiter.api.Test; + +public class TestMetadataAdapter { + + @Test + public void getAllReturnsEmptyIterableForMissingKey() { + MetadataAdapter headers = new MetadataAdapter(new Metadata()); + + assertNotNull(headers.getAll("missing")); + assertFalse(headers.getAll("missing").iterator().hasNext()); + assertNotNull(headers.getAllByte("missing-bin")); + assertFalse(headers.getAllByte("missing-bin").iterator().hasNext()); + } +} diff --git a/flight/flight-integration-tests/pom.xml b/flight/flight-integration-tests/pom.xml index 78a2d08ee1..f6ae8e16a5 100644 --- a/flight/flight-integration-tests/pom.xml +++ b/flight/flight-integration-tests/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-flight - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT flight-integration-tests @@ -58,7 +58,7 @@ under the License. commons-cli commons-cli - 1.9.0 + 1.11.0 org.slf4j @@ -101,7 +101,7 @@ under the License. - + META-INF/LICENSE.txt src/shade/LICENSE.txt diff --git a/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java b/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java index be746b5757..e400c031c2 100644 --- a/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java +++ b/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java @@ -98,6 +98,7 @@ static Schema getQuerySchema() { .isSearchable(true) .catalogName("catalog_test") .precision(100) + .remarks("test column") .build() .getMetadataMap()), null))); @@ -126,6 +127,7 @@ static Schema getQueryWithTransactionSchema() { .isSearchable(true) .catalogName("catalog_test") .precision(100) + .remarks("test column") .build() .getMetadataMap()), null))); diff --git a/flight/flight-integration-tests/src/shade/LICENSE.txt b/flight/flight-integration-tests/src/shade/LICENSE.txt index 484542fd20..3367c36f79 100644 --- a/flight/flight-integration-tests/src/shade/LICENSE.txt +++ b/flight/flight-integration-tests/src/shade/LICENSE.txt @@ -265,14 +265,14 @@ License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- -This binary artifact contains Guava 33.4.5-jre. +This binary artifact contains Guava 33.4.8-jre. Home page: https://github.com/google/guava License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- -This binary artifact contains Google Protobuf 4.30.1. +This binary artifact contains Google Protobuf 4.30.2. Copyright: Copyright 2008 Google Inc. All rights reserved. Home page: https://protobuf.dev/ diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml index 3362b60a8e..be2ee32868 100644 --- a/flight/flight-sql-jdbc-core/pom.xml +++ b/flight/flight-sql-jdbc-core/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-flight - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT flight-sql-jdbc-core @@ -105,7 +105,7 @@ under the License. commons-io commons-io - 2.18.0 + 2.22.0 test @@ -120,6 +120,31 @@ under the License. test + + com.squareup.okhttp3 + mockwebserver3 + 5.4.0 + test + + + com.squareup.okhttp3 + mockwebserver3-junit5 + 5.4.0 + test + + + com.squareup.okhttp3 + okhttp-jvm + 5.4.0 + test + + + com.squareup.okio + okio-jvm + 3.17.0 + test + + io.netty netty-common @@ -134,19 +159,32 @@ under the License. org.apache.calcite.avatica avatica - 1.26.0 + 1.27.0 org.bouncycastle bcpkix-jdk18on - 1.80 + 1.84 org.checkerframework checker-qual + + + com.github.ben-manes.caffeine + caffeine + 3.2.4 + + + + com.nimbusds + oauth2-oidc-sdk + 11.37.2 + + diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java index 3f072d071b..0110525fea 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadata.java @@ -45,6 +45,7 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.util.Arrays; +import java.util.Collections; import java.util.EnumMap; import java.util.HashMap; import java.util.List; @@ -75,18 +76,23 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaDatabaseMetaData; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Arrow Flight JDBC's implementation of {@link DatabaseMetaData}. */ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData { + private static final Logger LOGGER = LoggerFactory.getLogger(ArrowDatabaseMetadata.class); private static final String JAVA_REGEX_SPECIALS = "[]()|^-+*?{}$\\."; private static final Charset CHARSET = StandardCharsets.UTF_8; private static final byte[] EMPTY_BYTE_ARRAY = new byte[0]; @@ -164,6 +170,9 @@ public class ArrowDatabaseMetadata extends AvaticaDatabaseMetaData { LONGNVARCHAR, SqlSupportsConvert.SQL_CONVERT_LONGVARCHAR_VALUE); sqlTypesToFlightEnumConvertTypes.put(DATE, SqlSupportsConvert.SQL_CONVERT_DATE_VALUE); sqlTypesToFlightEnumConvertTypes.put(TIMESTAMP, SqlSupportsConvert.SQL_CONVERT_TIMESTAMP_VALUE); + + // Register the UUID extension type so it is always available for the driver + ExtensionTypeRegistry.register(UuidType.INSTANCE); } ArrowDatabaseMetadata(final AvaticaConnection connection) { @@ -769,7 +778,34 @@ private T getSqlInfoAndCacheIfCacheIsEmpty( } } } - return desiredType.cast(cachedSqlInfo.get(sqlInfoCommand)); + T value = desiredType.cast(cachedSqlInfo.get(sqlInfoCommand)); + if (value != null) { + return value; + } + LOGGER.debug( + "SqlInfo {} not provided by server, returning default for type {}", + sqlInfoCommand.name(), + desiredType.getSimpleName()); + + // Return sensible defaults when SqlInfo is unavailable + if (desiredType == Long.class) { + return desiredType.cast(0L); + } else if (desiredType == Integer.class) { + return desiredType.cast(0); + } else if (desiredType == Boolean.class) { + return desiredType.cast(false); + } else if (desiredType == String.class) { + return desiredType.cast(""); + } else if (desiredType == Map.class) { + return desiredType.cast(Collections.emptyMap()); + } else if (desiredType == List.class) { + return desiredType.cast(Collections.emptyList()); + } + + throw new SQLException( + String.format( + "The value of the SqlInfo %s is null and it could not be cast to %s.", + sqlInfoCommand.name(), desiredType.getName())); } private Optional convertListSqlInfoToString(final List sqlInfoList) { @@ -1066,6 +1102,7 @@ private int setGetColumnsVectorSchemaRootFromFields( (VarCharVector) currentRoot.getVector("IS_AUTOINCREMENT"); final VarCharVector isGeneratedColumnVector = (VarCharVector) currentRoot.getVector("IS_GENERATEDCOLUMN"); + final VarCharVector remarksVector = (VarCharVector) currentRoot.getVector("REMARKS"); for (int i = 0; i < tableColumnsSize; i++, ordinalIndex++) { final Field field = tableColumns.get(i); @@ -1139,6 +1176,11 @@ private int setGetColumnsVectorSchemaRootFromFields( isAutoincrementVector.setSafe(insertIndex, EMPTY_BYTE_ARRAY); } + String remarks = columnMetadata.getRemarks(); + if (remarks != null) { + remarksVector.setSafe(insertIndex, remarks.getBytes(CHARSET)); + } + // Fields also don't hold information about IS_AUTOINCREMENT and IS_GENERATEDCOLUMN, // so we're setting an empty string (as bytes), which means it couldn't be determined. isGeneratedColumnVector.setSafe(insertIndex, EMPTY_BYTE_ARRAY); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index cf9804d68b..623c2b81be 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -20,10 +20,14 @@ import io.netty.util.concurrent.DefaultThreadFactory; import java.sql.SQLException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; import java.util.Properties; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; +import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.memory.BufferAllocator; @@ -31,6 +35,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; +import org.apache.calcite.avatica.DriverVersion; /** Connection to the Arrow Flight server. */ public final class ArrowFlightConnection extends AvaticaConnection { @@ -39,6 +44,8 @@ public final class ArrowFlightConnection extends AvaticaConnection { private final ArrowFlightSqlClientHandler clientHandler; private final ArrowFlightConnectionConfigImpl config; private ExecutorService executorService; + private int metadataResultSetCount; + private Map metadataResultSetMap = new HashMap<>(); /** * Creates a new {@link ArrowFlightConnection}. @@ -63,6 +70,7 @@ private ArrowFlightConnection( this.config = Preconditions.checkNotNull(config, "Config cannot be null."); this.allocator = Preconditions.checkNotNull(allocator, "Allocator cannot be null."); this.clientHandler = Preconditions.checkNotNull(clientHandler, "Handler cannot be null."); + this.metadataResultSetCount = 0; } /** @@ -85,13 +93,16 @@ static ArrowFlightConnection createNewConnection( throws SQLException { url = replaceSemiColons(url); final ArrowFlightConnectionConfigImpl config = new ArrowFlightConnectionConfigImpl(properties); - final ArrowFlightSqlClientHandler clientHandler = createNewClientHandler(config, allocator); + final ArrowFlightSqlClientHandler clientHandler = + createNewClientHandler(config, allocator, driver.getDriverVersion()); return new ArrowFlightConnection( driver, factory, url, properties, config, allocator, clientHandler); } private static ArrowFlightSqlClientHandler createNewClientHandler( - final ArrowFlightConnectionConfigImpl config, final BufferAllocator allocator) + final ArrowFlightConnectionConfigImpl config, + final BufferAllocator allocator, + final DriverVersion driverVersion) throws SQLException { try { return new ArrowFlightSqlClientHandler.Builder() @@ -113,7 +124,10 @@ private static ArrowFlightSqlClientHandler createNewClientHandler( .withRetainCookies(config.retainCookies()) .withRetainAuth(config.retainAuth()) .withCatalog(config.getCatalog()) + .withClientCache(config.useClientCache() ? new FlightClientCache() : null) .withConnectTimeout(config.getConnectTimeout()) + .withDriverVersion(driverVersion) + .withOAuthConfiguration(config.getOauthConfiguration()) .build(); } catch (final SQLException e) { try { @@ -164,6 +178,31 @@ synchronized ExecutorService getExecutorService() { : executorService; } + /** + * Registers a new metadata ResultSet and assigns it a unique ID. Metadata ResultSets are those + * created without an associated Statement. + * + * @param resultSet the ResultSet to register + * @return the assigned ID + */ + int getNewMetadataResultSetId(ArrowFlightJdbcFlightStreamResultSet resultSet) { + metadataResultSetMap.put(metadataResultSetCount, resultSet); + return metadataResultSetCount++; + } + + /** + * Unregisters a metadata ResultSet when it is closed. This method is called by metadata + * ResultSets during their close operation to remove themselves from the tracking map. + * + * @param id the ID of the ResultSet to unregister, or null if not a metadata ResultSet + */ + void onResultSetClose(Integer id) { + if (id == null) { + return; + } + metadataResultSetMap.remove(id); + } + @Override public Properties getClientInfo() { final Properties copy = new Properties(); @@ -173,19 +212,41 @@ public Properties getClientInfo() { @Override public void close() throws SQLException { - clientHandler.close(); - if (executorService != null) { - executorService.shutdown(); + Exception topLevelException = null; + try { + if (executorService != null) { + executorService.shutdown(); + } + } catch (final Exception e) { + topLevelException = e; + } + // copies of the collections are used to avoid concurrent modification problems + ArrayList closeables = new ArrayList<>(statementMap.values()); + closeables.addAll(new ArrayList<>(metadataResultSetMap.values())); + closeables.add(clientHandler); + closeables.addAll(allocator.getChildAllocators()); + closeables.add(allocator); + try { + AutoCloseables.close(closeables); + } catch (final Exception e) { + if (topLevelException == null) { + topLevelException = e; + } else { + topLevelException.addSuppressed(e); + } } - try { - AutoCloseables.close(clientHandler); - allocator.getChildAllocators().forEach(AutoCloseables::closeNoChecked); - AutoCloseables.close(allocator); - super.close(); } catch (final Exception e) { - throw AvaticaConnection.HELPER.createException(e.getMessage(), e); + if (topLevelException == null) { + topLevelException = e; + } else { + topLevelException.addSuppressed(e); + } + } + if (topLevelException != null) { + throw AvaticaConnection.HELPER.createException( + topLevelException.getMessage(), topLevelException); } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java index 9b9eba51e5..f3d76ace92 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArray.java @@ -26,6 +26,7 @@ import org.apache.arrow.driver.jdbc.utils.SqlTypes; import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -135,12 +136,22 @@ public ResultSet getResultSet(long index, int count) throws SQLException { private static ResultSet getResultSetNoBoundariesCheck( ValueVector dataVector, long start, long count) throws SQLException { + int intStart = LargeMemoryUtil.checkedCastToInt(start); + int intCount = LargeMemoryUtil.checkedCastToInt(count); + + // Create an index vector with 1-based indices (per JDBC spec) to return with value vector + IntVector indexVector = new IntVector("INDEX", dataVector.getAllocator()); + indexVector.allocateNew(intCount); + for (int i = 0; i < intCount; i++) { + indexVector.set(i, i + 1); + } + indexVector.setValueCount(intCount); + TransferPair transferPair = dataVector.getTransferPair(dataVector.getAllocator()); - transferPair.splitAndTransfer( - LargeMemoryUtil.checkedCastToInt(start), LargeMemoryUtil.checkedCastToInt(count)); - FieldVector vectorSlice = (FieldVector) transferPair.getTo(); + transferPair.splitAndTransfer(intStart, intCount); + FieldVector valueVector = (FieldVector) transferPair.getTo(); - VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(vectorSlice); + VectorSchemaRoot vectorSchemaRoot = VectorSchemaRoot.of(indexVector, valueVector); return ArrowFlightJdbcVectorSchemaRootResultSet.fromVectorSchemaRoot(vectorSchemaRoot); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java index 53e6120f62..12ef8030d7 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriver.java @@ -75,7 +75,9 @@ public Logger getParentLogger() { public ArrowFlightConnection connect(final String url, final Properties info) throws SQLException { final Properties properties = new Properties(info); - properties.putAll(info); + if (info != null) { + properties.putAll(info); + } if (url != null) { final Optional> maybeProperties = getUrlsArgs(url); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java index aabaf01e63..376e5b11e7 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java @@ -54,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet private VectorSchemaRoot currentVectorSchemaRoot; private Schema schema; + private Integer id = null; // used for metadata result sets only /** Public constructor used by ArrowFlightJdbcFactory. */ ArrowFlightJdbcFlightStreamResultSet( @@ -82,6 +83,7 @@ private ArrowFlightJdbcFlightStreamResultSet( super(null, state, signature, resultSetMetaData, timeZone, firstFrame); this.connection = connection; this.flightInfo = flightInfo; + this.id = connection.getNewMetadataResultSetId(this); } /** @@ -104,7 +106,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo( final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); - final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); + final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); @@ -234,7 +236,12 @@ protected void cancel() { @Override public synchronized void close() { + try { + if (isClosed()) { + return; + } + this.connection.onResultSetClose(id); if (flightEndpointDataQueue != null) { // flightStreamQueue should close currentFlightStream internally flightEndpointDataQueue.close(); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java index 0dc2b07c97..ad6670a001 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java @@ -19,23 +19,26 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.Types; import java.util.HashSet; import java.util.List; -import java.util.Objects; import java.util.Set; import java.util.TimeZone; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaResultSetMetaData; +import org.apache.calcite.avatica.AvaticaSite; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.Meta; import org.apache.calcite.avatica.Meta.Frame; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.QueryState; +import org.apache.calcite.avatica.util.Cursor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -70,7 +73,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot( final TimeZone timeZone = TimeZone.getDefault(); final QueryState state = new QueryState(); - final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null); + final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null); final AvaticaResultSetMetaData resultSetMetaData = new AvaticaResultSetMetaData(null, null, signature); @@ -102,6 +105,33 @@ void populateData(final VectorSchemaRoot vectorSchemaRoot, final Schema schema) execute2(new ArrowFlightJdbcCursor(vectorSchemaRoot), this.signature.columns); } + /** + * The default method in AvaticaResultSet does not properly handle TIMESTASMP_WITH_TIMEZONE, so we + * override here to add support. + * + * @param columnIndex the first column is 1, the second is 2, ... + * @return Object + * @throws SQLException if there is an underlying exception + */ + @Override + public Object getObject(int columnIndex) throws SQLException { + this.checkOpen(); + + Cursor.Accessor accessor; + try { + accessor = accessorList.get(columnIndex - 1); + } catch (IndexOutOfBoundsException e) { + throw AvaticaConnection.HELPER.createException("invalid column ordinal: " + columnIndex); + } + + ColumnMetaData metaData = columnMetaDataList.get(columnIndex - 1); + if (metaData.type.id == Types.TIMESTAMP_WITH_TIMEZONE) { + return accessor.getTimestamp(localCalendar); + } else { + return AvaticaSite.get(accessor, metaData.type.id, true, localCalendar); + } + } + @Override protected void cancel() { signature.columns.clear(); @@ -128,12 +158,10 @@ public void close() { } catch (final Exception e) { exceptions.add(e); } - if (!Objects.isNull(statement)) { - try { - super.close(); - } catch (final Exception e) { - exceptions.add(e); - } + try { + super.close(); + } catch (final Exception e) { + exceptions.add(e); } exceptions.parallelStream().forEach(e -> LOGGER.error(e.getMessage(), e)); exceptions.stream() diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java index 9c7112f1c3..0d85b5eddb 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java @@ -53,7 +53,8 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) { } /** Construct a signature. */ - static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) { + static Signature newSignature( + final String sql, Schema resultSetSchema, Schema parameterSchema, Boolean isUpdate) { List columnMetaData = resultSetSchema == null ? new ArrayList<>() @@ -62,21 +63,32 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p parameterSchema == null ? new ArrayList<>() : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); - + // If the server provided the is_update field, use it to determine the statement type + StatementType statementType; + if (isUpdate != null) { + statementType = isUpdate ? StatementType.IS_DML : StatementType.SELECT; + } else { + // Fall back to the legacy logic: check if the result set schema is empty + statementType = + resultSetSchema == null || resultSetSchema.getFields().isEmpty() + ? StatementType.IS_DML + : StatementType.SELECT; + } return new Signature( columnMetaData, sql, parameters, Collections.emptyMap(), null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor - StatementType.SELECT); + statementType); } @Override public void closeStatement(final StatementHandle statementHandle) { PreparedStatement preparedStatement = statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle)); - // Testing if the prepared statement was created because the statement can be not created until + // Testing if the prepared statement was created because the statement can be + // not created until // this moment if (preparedStatement != null) { preparedStatement.close(); @@ -105,7 +117,8 @@ public ExecuteResult execute( preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) .bind(typedValues); - if (statementHandle.signature == null) { + if (statementHandle.signature == null + || statementHandle.signature.statementType == StatementType.IS_DML) { // Update query long updatedCount = preparedStatement.executeUpdate(); return new ExecuteResult( @@ -173,7 +186,10 @@ private PreparedStatement prepareForHandle(final String query, StatementHandle h ((ArrowFlightConnection) connection).getClientHandler().prepare(query); handle.signature = newSignature( - query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); + query, + preparedStatement.getDataSetSchema(), + preparedStatement.getParameterSchema(), + preparedStatement.isUpdate()); statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); return preparedStatement; } @@ -220,7 +236,8 @@ public ExecuteResult prepareAndExecute( MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null); return new ExecuteResult(Collections.singletonList(metaResultSet)); } catch (SQLTimeoutException e) { - // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and Runtime + // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and + // Runtime // Exceptions. throw new RuntimeException(e); } catch (SQLException e) { @@ -249,6 +266,20 @@ public boolean syncResults( return false; } + @Override + public ConnectionProperties connectionSync(ConnectionHandle ch, ConnectionProperties connProps) { + final ConnectionProperties result = super.connectionSync(ch, connProps); + final String newCatalog = this.connProps.getCatalog(); + if (newCatalog != null) { + try { + ((ArrowFlightConnection) connection).getClientHandler().setCatalog(newCatalog); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + return result; + } + void setDefaultConnectionProperties() { // TODO Double-check this. connProps @@ -264,7 +295,8 @@ PreparedStatement getPreparedStatement(StatementHandle statementHandle) { return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle)); } - // Helper used to look up prepared statement instances later. Avatica doesn't give us the + // Helper used to look up prepared statement instances later. Avatica doesn't + // give us the // signature in // an UPDATE code path so we can't directly use StatementHandle as a map key. private static final class StatementHandleKey { diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java index f0fa55fa82..cd762fb1ac 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessor.java @@ -36,6 +36,10 @@ import java.util.Calendar; import java.util.Map; import java.util.function.IntSupplier; +import org.joou.UByte; +import org.joou.UInteger; +import org.joou.ULong; +import org.joou.UShort; /** Base Jdbc Accessor. */ public abstract class ArrowFlightJdbcAccessor implements Accessor { @@ -99,6 +103,26 @@ public long getLong() throws SQLException { throw getOperationNotSupported(this.getClass()); } + @Override + public UByte getUByte() throws SQLException { + throw getOperationNotSupported(this.getClass()); + } + + @Override + public UShort getUShort() throws SQLException { + throw getOperationNotSupported(this.getClass()); + } + + @Override + public UInteger getUInt() throws SQLException { + throw getOperationNotSupported(this.getClass()); + } + + @Override + public ULong getULong() throws SQLException { + throw getOperationNotSupported(this.getClass()); + } + @Override public float getFloat() throws SQLException { throw getOperationNotSupported(this.getClass()); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java index dad1fa5f73..8362eb7627 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactory.java @@ -19,6 +19,7 @@ import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.impl.ArrowFlightJdbcNullVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor; +import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor; @@ -65,9 +66,12 @@ import org.apache.arrow.vector.UInt2Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarBinaryVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; @@ -130,9 +134,15 @@ public static ArrowFlightJdbcAccessor createAccessor( } else if (vector instanceof VarBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( (VarBinaryVector) vector, getCurrentRow, setCursorWasNull); + } else if (vector instanceof ViewVarBinaryVector) { + return new ArrowFlightJdbcBinaryVectorAccessor( + (ViewVarBinaryVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof LargeVarBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( (LargeVarBinaryVector) vector, getCurrentRow, setCursorWasNull); + } else if (vector instanceof UuidVector) { + return new ArrowFlightJdbcUuidVectorAccessor( + (UuidVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof FixedSizeBinaryVector) { return new ArrowFlightJdbcBinaryVectorAccessor( (FixedSizeBinaryVector) vector, getCurrentRow, setCursorWasNull); @@ -163,6 +173,9 @@ public static ArrowFlightJdbcAccessor createAccessor( } else if (vector instanceof LargeVarCharVector) { return new ArrowFlightJdbcVarCharVectorAccessor( (LargeVarCharVector) vector, getCurrentRow, setCursorWasNull); + } else if (vector instanceof ViewVarCharVector) { + return new ArrowFlightJdbcVarCharVectorAccessor( + (ViewVarCharVector) vector, getCurrentRow, setCursorWasNull); } else if (vector instanceof DurationVector) { return new ArrowFlightJdbcDurationVectorAccessor( (DurationVector) vector, getCurrentRow, setCursorWasNull); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java index 30dfffce64..e71b6380a9 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcBinaryVectorAccessor.java @@ -27,6 +27,7 @@ import org.apache.arrow.vector.FixedSizeBinaryVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.ViewVarBinaryVector; /** * Accessor for the Arrow types: {@link FixedSizeBinaryVector}, {@link VarBinaryVector} and {@link @@ -61,6 +62,13 @@ public ArrowFlightJdbcBinaryVectorAccessor( this(vector::get, currentRowSupplier, setCursorWasNull); } + public ArrowFlightJdbcBinaryVectorAccessor( + ViewVarBinaryVector vector, + IntSupplier currentRowSupplier, + ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { + this(vector::get, currentRowSupplier, setCursorWasNull); + } + private ArrowFlightJdbcBinaryVectorAccessor( ByteArrayGetter getter, IntSupplier currentRowSupplier, diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java new file mode 100644 index 0000000000..4bdbcbb63d --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessor.java @@ -0,0 +1,88 @@ +/* + * 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.arrow.driver.jdbc.accessor.impl.binary; + +import java.util.UUID; +import java.util.function.IntSupplier; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessor; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.util.UuidUtility; + +/** + * Accessor for the Arrow UUID extension type ({@link UuidVector}). + * + *

This accessor provides JDBC-compatible access to UUID values stored in Arrow's canonical UUID + * extension type ('arrow.uuid'). It follows PostgreSQL JDBC driver conventions: + * + *

    + *
  • {@link #getObject()} returns {@link java.util.UUID} + *
  • {@link #getString()} returns the hyphenated string format (e.g., + * "550e8400-e29b-41d4-a716-446655440000") + *
  • {@link #getBytes()} returns the 16-byte binary representation + *
+ */ +public class ArrowFlightJdbcUuidVectorAccessor extends ArrowFlightJdbcAccessor { + + private final UuidVector vector; + + /** + * Creates a new accessor for a UUID vector. + * + * @param vector the UUID vector to access + * @param currentRowSupplier supplier for the current row index + * @param setCursorWasNull consumer to set the wasNull flag + */ + public ArrowFlightJdbcUuidVectorAccessor( + UuidVector vector, + IntSupplier currentRowSupplier, + ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { + super(currentRowSupplier, setCursorWasNull); + this.vector = vector; + } + + @Override + public Object getObject() { + UUID uuid = vector.getObject(getCurrentRow()); + this.wasNull = uuid == null; + this.wasNullConsumer.setWasNull(this.wasNull); + return uuid; + } + + @Override + public Class getObjectClass() { + return UUID.class; + } + + @Override + public String getString() { + UUID uuid = (UUID) getObject(); + if (uuid == null) { + return null; + } + return uuid.toString(); + } + + @Override + public byte[] getBytes() { + UUID uuid = (UUID) getObject(); + if (uuid == null) { + return null; + } + return UuidUtility.getBytesFromUUID(uuid); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java index ebebf6ca74..7b04e89346 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessor.java @@ -35,6 +35,7 @@ import org.apache.arrow.driver.jdbc.utils.DateTimeUtils; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.util.Text; /** Accessor for the Arrow types: {@link VarCharVector} and {@link LargeVarCharVector}. */ @@ -62,6 +63,13 @@ public ArrowFlightJdbcVarCharVectorAccessor( this(vector::get, currentRowSupplier, setCursorWasNull); } + public ArrowFlightJdbcVarCharVectorAccessor( + ViewVarCharVector vector, + IntSupplier currentRowSupplier, + ArrowFlightJdbcAccessorFactory.WasNullConsumer setCursorWasNull) { + this(vector::get, currentRowSupplier, setCursorWasNull); + } + ArrowFlightJdbcVarCharVectorAccessor( Getter getter, IntSupplier currentRowSupplier, diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java index cbbe223eb8..719cc38a2b 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java @@ -32,7 +32,12 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration; +import org.apache.arrow.driver.jdbc.client.oauth.OAuthCredentialWriter; +import org.apache.arrow.driver.jdbc.client.oauth.OAuthTokenProvider; import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils; +import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache; +import org.apache.arrow.driver.jdbc.client.utils.FlightLocationQueue; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CloseSessionRequest; @@ -45,7 +50,6 @@ import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.LocationSchemes; -import org.apache.arrow.flight.SessionOptionValue; import org.apache.arrow.flight.SessionOptionValueFactory; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; @@ -64,6 +68,7 @@ import org.apache.arrow.util.VisibleForTesting; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.DriverVersion; import org.apache.calcite.avatica.Meta.StatementType; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; @@ -75,21 +80,27 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable { // JDBC connection string query parameter private static final String CATALOG = "catalog"; + private final String cacheKey; private final FlightSqlClient sqlClient; private final Set options = new HashSet<>(); private final Builder builder; private final Optional catalog; + private final @Nullable FlightClientCache flightClientCache; ArrowFlightSqlClientHandler( + final String cacheKey, final FlightSqlClient sqlClient, final Builder builder, final Collection credentialOptions, - final Optional catalog) { + final Optional catalog, + final @Nullable FlightClientCache flightClientCache) { this.options.addAll(builder.options); this.options.addAll(credentialOptions); + this.cacheKey = Preconditions.checkNotNull(cacheKey); this.sqlClient = Preconditions.checkNotNull(sqlClient); this.builder = builder; this.catalog = catalog; + this.flightClientCache = flightClientCache; } /** @@ -101,12 +112,15 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable { * @return a new {@link ArrowFlightSqlClientHandler}. */ static ArrowFlightSqlClientHandler createNewHandler( + final String cacheKey, final FlightClient client, final Builder builder, final Collection options, - final Optional catalog) { + final Optional catalog, + final @Nullable FlightClientCache flightClientCache) { final ArrowFlightSqlClientHandler handler = - new ArrowFlightSqlClientHandler(new FlightSqlClient(client), builder, options, catalog); + new ArrowFlightSqlClientHandler( + cacheKey, new FlightSqlClient(client), builder, options, catalog, flightClientCache); handler.setSetCatalogInSessionIfPresent(); return handler; } @@ -135,22 +149,33 @@ public List getStreams(final FlightInfo flightInfo) try { for (FlightEndpoint endpoint : flightInfo.getEndpoints()) { if (endpoint.getLocations().isEmpty()) { - // Create a stream using the current client only and do not close the client at the end. + // Create a stream using the current client only and do not close the client at + // the end. endpoints.add( new CloseableEndpointStreamPair( sqlClient.getStream(endpoint.getTicket(), getOptions()), null)); } else { // Clone the builder and then set the new endpoint on it. - // GH-38574: Currently a new FlightClient will be made for each partition that returns a - // non-empty Location then disposed of. It may be better to cache clients because a server - // may report the same Locations. It would also be good to identify when the reported + // GH-38574: Currently a new FlightClient will be made for each partition that + // returns a + // non-empty Location then disposed of. It may be better to cache clients + // because a server + // may report the same Locations. It would also be good to identify when the + // reported // location - // is the same as the original connection's Location and skip creating a FlightClient in + // is the same as the original connection's Location and skip creating a + // FlightClient in // that scenario. + // Also copy the cache to the client so we can share a cache. Cache needs to + // cache + // negative attempts too. List exceptions = new ArrayList<>(); CloseableEndpointStreamPair stream = null; - for (Location location : endpoint.getLocations()) { + FlightLocationQueue locations = + new FlightLocationQueue(flightClientCache, endpoint.getLocations()); + while (locations.hasNext()) { + Location location = locations.next(); final URI endpointUri = location.getUri(); if (endpointUri.getScheme().equals(LocationSchemes.REUSE_CONNECTION)) { stream = @@ -163,6 +188,7 @@ public List getStreams(final FlightInfo flightInfo) .withHost(endpointUri.getHost()) .withPort(endpointUri.getPort()) .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS)) + .withClientCache(flightClientCache) .withConnectTimeout(builder.connectTimeout); ArrowFlightSqlClientHandler endpointHandler = null; @@ -177,12 +203,29 @@ public List getStreams(final FlightInfo flightInfo) stream.getStream().getSchema(); } catch (Exception ex) { if (endpointHandler != null) { + // If the exception is related to connectivity, mark the client as a dud. + if (flightClientCache != null) { + if (ex instanceof FlightRuntimeException + && ((FlightRuntimeException) ex).status().code() + == FlightStatusCode.UNAVAILABLE + && + // IOException covers SocketException and Netty's (private) + // AnnotatedSocketException + // We are looking for things like "Network is unreachable" + ex.getCause() instanceof IOException) { + flightClientCache.markLocationAsDud(location.toString()); + } + } + AutoCloseables.close(endpointHandler); } exceptions.add(ex); continue; } + if (flightClientCache != null) { + flightClientCache.markLocationAsReachable(location.toString()); + } break; } if (stream != null) { @@ -227,15 +270,86 @@ public FlightInfo getInfo(final String query) { @Override public void close() throws SQLException { if (catalog.isPresent()) { - sqlClient.closeSession(new CloseSessionRequest(), getOptions()); + try { + sqlClient.closeSession(new CloseSessionRequest(), getOptions()); + } catch (FlightRuntimeException fre) { + handleBenignCloseException( + fre, "Failed to close Flight SQL session.", "closing Flight SQL session"); + } } try { AutoCloseables.close(sqlClient); + } catch (FlightRuntimeException fre) { + handleBenignCloseException( + fre, "Failed to clean up client resources.", "closing Flight SQL client"); } catch (final Exception e) { throw new SQLException("Failed to clean up client resources.", e); } } + /** + * Handles FlightRuntimeException during close operations, suppressing benign gRPC shutdown errors + * while re-throwing genuine failures. + * + * @param fre the FlightRuntimeException to handle + * @param sqlErrorMessage the SQLException message to use for genuine failures + * @param operationDescription description of the operation for logging + * @throws SQLException if the exception represents a genuine failure + */ + private void handleBenignCloseException( + FlightRuntimeException fre, String sqlErrorMessage, String operationDescription) + throws SQLException { + if (isBenignCloseException(fre)) { + logSuppressedCloseException(fre, operationDescription); + } else { + throw new SQLException(sqlErrorMessage, fre); + } + } + + /** + * Handles FlightRuntimeException during close operations, suppressing benign gRPC shutdown errors + * while re-throwing genuine failures as FlightRuntimeException. + * + * @param fre the FlightRuntimeException to handle + * @param operationDescription description of the operation for logging + * @throws FlightRuntimeException if the exception represents a genuine failure + */ + private void handleBenignCloseException(FlightRuntimeException fre, String operationDescription) + throws FlightRuntimeException { + if (isBenignCloseException(fre)) { + logSuppressedCloseException(fre, operationDescription); + } else { + throw fre; + } + } + + /** + * Determines if a FlightRuntimeException represents a benign close operation error that should be + * suppressed. + * + * @param fre the FlightRuntimeException to check + * @return true if the exception should be suppressed, false otherwise + */ + private boolean isBenignCloseException(FlightRuntimeException fre) { + return fre.status().code().equals(FlightStatusCode.UNAVAILABLE) + || (fre.status().code().equals(FlightStatusCode.INTERNAL) + && fre.getMessage() != null + && fre.getMessage().contains("Connection closed after GOAWAY")); + } + + /** + * Logs a suppressed close exception with appropriate level based on debug settings. + * + * @param fre the FlightRuntimeException being suppressed + * @param operationDescription description of the operation for logging + */ + private void logSuppressedCloseException( + FlightRuntimeException fre, String operationDescription) { + // ARROW-17785 and GH-863: suppress exceptions caused by flaky gRPC layer during + // shutdown + LOGGER.debug("Suppressed error {}", operationDescription, fre); + } + /** A prepared statement handler. */ public interface PreparedStatement extends AutoCloseable { /** @@ -274,6 +388,14 @@ public interface PreparedStatement extends AutoCloseable { */ Schema getParameterSchema(); + /** + * Gets whether this {@link PreparedStatement} is an update statement. + * + * @return {@code true} if this is an update statement, {@code false} if it's a query, or {@code + * null} if the server did not provide this information. + */ + @Nullable Boolean isUpdate(); + void setParameters(VectorSchemaRoot parameters); @Override @@ -283,25 +405,40 @@ public interface PreparedStatement extends AutoCloseable { /** A connection is created with catalog set as a session option. */ private void setSetCatalogInSessionIfPresent() { if (catalog.isPresent()) { - final SetSessionOptionsRequest setSessionOptionRequest = - new SetSessionOptionsRequest( - ImmutableMap.builder() - .put(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog.get())) - .build()); - final SetSessionOptionsResult result = - sqlClient.setSessionOptions(setSessionOptionRequest, getOptions()); + try { + setCatalog(catalog.get()); + } catch (SQLException e) { + throw CallStatus.INVALID_ARGUMENT + .withDescription(e.getMessage()) + .withCause(e) + .toRuntimeException(); + } + } + } + /** + * Sets the catalog for the current session. + * + * @param catalog the catalog to set. + * @throws SQLException if an error occurs while setting the catalog. + */ + public void setCatalog(final String catalog) throws SQLException { + final SetSessionOptionsRequest request = + new SetSessionOptionsRequest( + ImmutableMap.of(CATALOG, SessionOptionValueFactory.makeSessionOptionValue(catalog))); + try { + final SetSessionOptionsResult result = sqlClient.setSessionOptions(request, getOptions()); if (result.hasErrors()) { - Map errors = result.getErrors(); - for (Map.Entry error : errors.entrySet()) { + final Map errors = result.getErrors(); + for (final Map.Entry error : errors.entrySet()) { LOGGER.warn(error.toString()); } - throw CallStatus.INVALID_ARGUMENT - .withDescription( - String.format( - "Cannot set session option for catalog = %s. Check log for details.", catalog)) - .toRuntimeException(); + throw new SQLException( + String.format( + "Cannot set session option for catalog = %s. Check log for details.", catalog)); } + } catch (final FlightRuntimeException e) { + throw new SQLException(e); } } @@ -327,6 +464,12 @@ public long executeUpdate() { @Override public StatementType getType() { + // If the server provided the is_update field, use it to determine the statement type + final Boolean isUpdate = preparedStatement.isUpdate(); + if (isUpdate != null) { + return isUpdate ? StatementType.UPDATE : StatementType.SELECT; + } + // Fall back to the legacy logic: check if the result set schema is empty final Schema schema = preparedStatement.getResultSetSchema(); return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT; } @@ -346,19 +489,17 @@ public void setParameters(VectorSchemaRoot parameters) { preparedStatement.setParameters(parameters); } + @Override + public Boolean isUpdate() { + return preparedStatement.isUpdate(); + } + @Override public void close() { try { preparedStatement.close(getOptions()); } catch (FlightRuntimeException fre) { - // ARROW-17785: suppress exceptions caused by flaky gRPC layer - if (fre.status().code().equals(FlightStatusCode.UNAVAILABLE) - || (fre.status().code().equals(FlightStatusCode.INTERNAL) - && fre.getMessage().contains("Connection closed after GOAWAY"))) { - LOGGER.warn("Supressed error closing PreparedStatement", fre); - return; - } - throw fre; + handleBenignCloseException(fre, "closing PreparedStatement"); } } }; @@ -514,6 +655,9 @@ public FlightInfo getCrossReference( /** Builder for {@link ArrowFlightSqlClientHandler}. */ public static final class Builder { + static final String USER_AGENT_TEMPLATE = "JDBC Flight SQL Driver %s"; + static final String DEFAULT_VERSION = "(unknown or development build)"; + private final Set middlewareFactories = new HashSet<>(); private final Set options = new HashSet<>(); private String host; @@ -549,9 +693,14 @@ public static final class Builder { @VisibleForTesting Optional catalog = Optional.empty(); + @VisibleForTesting @Nullable FlightClientCache flightClientCache; + @VisibleForTesting @Nullable Duration connectTimeout; - // These two middleware are for internal use within build() and should not be exposed by builder + @VisibleForTesting @Nullable OAuthConfiguration oauthConfig; + + // These two middleware are for internal use within build() and should not be + // exposed by builder // APIs. // Note that these middleware may not necessarily be registered. @VisibleForTesting @@ -561,6 +710,8 @@ public static final class Builder { @VisibleForTesting ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory(); + DriverVersion driverVersion; + public Builder() {} /** @@ -587,6 +738,7 @@ public Builder() {} this.clientKeyPath = original.clientKeyPath; this.allocator = original.allocator; this.catalog = original.catalog; + this.oauthConfig = original.oauthConfig; if (original.retainCookies) { this.cookieFactory = original.cookieFactory; @@ -595,6 +747,8 @@ public Builder() {} if (original.retainAuth) { this.authFactory = original.authFactory; } + + this.driverVersion = original.driverVersion; } /** @@ -833,11 +987,42 @@ public Builder withCatalog(@Nullable final String catalog) { return this; } + public Builder withClientCache(FlightClientCache flightClientCache) { + this.flightClientCache = flightClientCache; + return this; + } + public Builder withConnectTimeout(Duration connectTimeout) { this.connectTimeout = connectTimeout; return this; } + /** + * Sets the driver version for this handler. + * + * @param driverVersion the driver version to set + * @return this builder instance + */ + public Builder withDriverVersion(DriverVersion driverVersion) { + this.driverVersion = driverVersion; + return this; + } + + /** + * Sets the OAuth configuration for this handler. + * + * @param oauthConfig the OAuth configuration + * @return this builder instance + */ + public Builder withOAuthConfiguration(final OAuthConfiguration oauthConfig) { + this.oauthConfig = oauthConfig; + return this; + } + + public String getCacheKey() { + return getLocation().toString(); + } + /** Get the location that this client will connect to. */ public Location getLocation() { if (useEncryption) { @@ -853,7 +1038,8 @@ public Location getLocation() { * @throws SQLException on error. */ public ArrowFlightSqlClientHandler build() throws SQLException { - // Copy middleware so that the build method doesn't change the state of the builder fields + // Copy middleware so that the build method doesn't change the state of the + // builder fields // itself. Set buildTimeMiddlewareFactories = new HashSet<>(this.middlewareFactories); @@ -861,7 +1047,8 @@ public ArrowFlightSqlClientHandler build() throws SQLException { boolean isUsingUserPasswordAuth = username != null && token == null; try { - // Token should take priority since some apps pass in a username/password even when a token + // Token should take priority since some apps pass in a username/password even + // when a token // is provided if (isUsingUserPasswordAuth) { buildTimeMiddlewareFactories.add(authFactory); @@ -869,6 +1056,11 @@ public ArrowFlightSqlClientHandler build() throws SQLException { final NettyClientBuilder clientBuilder = new NettyClientBuilder(); clientBuilder.allocator(allocator); + String userAgent = String.format(USER_AGENT_TEMPLATE, DEFAULT_VERSION); + if (driverVersion != null && driverVersion.versionString != null) { + userAgent = String.format(USER_AGENT_TEMPLATE, driverVersion.versionString); + } + buildTimeMiddlewareFactories.add(new ClientCookieMiddleware.Factory()); buildTimeMiddlewareFactories.forEach(clientBuilder::intercept); if (useEncryption) { @@ -903,6 +1095,9 @@ public ArrowFlightSqlClientHandler build() throws SQLException { } NettyChannelBuilder channelBuilder = clientBuilder.build(); + + channelBuilder.userAgent(userAgent); + if (connectTimeout != null) { channelBuilder.withOption( ChannelOption.CONNECT_TIMEOUT_MILLIS, (int) connectTimeout.toMillis()); @@ -911,9 +1106,15 @@ public ArrowFlightSqlClientHandler build() throws SQLException { FlightGrpcUtils.createFlightClient( allocator, channelBuilder.build(), clientBuilder.middleware()); final ArrayList credentialOptions = new ArrayList<>(); - if (isUsingUserPasswordAuth) { - // If the authFactory has already been used for a handshake, use the existing token. - // This can occur if the authFactory is being re-used for a new connection spawned for + // Authentication priority: OAuth > token > username/password + if (oauthConfig != null) { + OAuthTokenProvider tokenProvider = oauthConfig.createTokenProvider(); + credentialOptions.add(new CredentialCallOption(new OAuthCredentialWriter(tokenProvider))); + } else if (isUsingUserPasswordAuth) { + // If the authFactory has already been used for a handshake, use the existing + // token. + // This can occur if the authFactory is being re-used for a new connection + // spawned for // getStream(). if (authFactory.getCredentialCallOption() != null) { credentialOptions.add(authFactory.getCredentialCallOption()); @@ -931,7 +1132,7 @@ public ArrowFlightSqlClientHandler build() throws SQLException { options.toArray(new CallOption[0]))); } return ArrowFlightSqlClientHandler.createNewHandler( - client, this, credentialOptions, catalog); + getCacheKey(), client, this, credentialOptions, catalog, flightClientCache); } catch (final IllegalArgumentException | GeneralSecurityException diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java new file mode 100644 index 0000000000..9c377a5850 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/AbstractOAuthTokenProvider.java @@ -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.arrow.driver.jdbc.client.oauth; + +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.Scope; +import com.nimbusds.oauth2.sdk.TokenErrorResponse; +import com.nimbusds.oauth2.sdk.TokenRequest; +import com.nimbusds.oauth2.sdk.TokenResponse; +import com.nimbusds.oauth2.sdk.auth.ClientAuthentication; +import com.nimbusds.oauth2.sdk.token.AccessToken; +import java.io.IOException; +import java.net.URI; +import java.sql.SQLException; +import java.time.Instant; +import org.apache.arrow.util.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Abstract base class for OAuth token providers that handles token caching, refresh logic, and + * common request/response handling. + */ +public abstract class AbstractOAuthTokenProvider implements OAuthTokenProvider { + protected static final int EXPIRATION_BUFFER_SECONDS = 30; + protected static final int DEFAULT_EXPIRATION_SECONDS = 3600; + + private final Object tokenLock = new Object(); + private volatile @Nullable TokenInfo cachedToken; + + @VisibleForTesting URI tokenUri; + + @VisibleForTesting @Nullable ClientAuthentication clientAuth; + + @VisibleForTesting @Nullable Scope scope; + + @Override + public String getValidToken() throws SQLException { + TokenInfo token = cachedToken; + if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) { + return token.getAccessToken(); + } + + synchronized (tokenLock) { + token = cachedToken; + if (token != null && !token.isExpired(EXPIRATION_BUFFER_SECONDS)) { + return token.getAccessToken(); + } + cachedToken = fetchNewToken(); + return cachedToken.getAccessToken(); + } + } + + /** + * Fetches a new token from the authorization server. This method handles the common + * request/response logic while delegating flow-specific request building to subclasses. + * + * @return the new token information + * @throws SQLException if token cannot be obtained + */ + protected TokenInfo fetchNewToken() throws SQLException { + try { + TokenRequest request = buildTokenRequest(); + TokenResponse response = TokenResponse.parse(request.toHTTPRequest().send()); + + if (!response.indicatesSuccess()) { + TokenErrorResponse errorResponse = response.toErrorResponse(); + String errorMsg = + String.format( + "OAuth request failed: %s - %s", + errorResponse.getErrorObject().getCode(), + errorResponse.getErrorObject().getDescription()); + throw new SQLException(errorMsg); + } + + AccessToken accessToken = response.toSuccessResponse().getTokens().getAccessToken(); + long expiresIn = + accessToken.getLifetime() > 0 ? accessToken.getLifetime() : DEFAULT_EXPIRATION_SECONDS; + Instant expiresAt = Instant.now().plusSeconds(expiresIn); + + return new TokenInfo(accessToken.getValue(), expiresAt); + } catch (ParseException e) { + throw new SQLException("Failed to parse OAuth token response", e); + } catch (IOException e) { + throw new SQLException("Failed to send OAuth token request", e); + } + } + + /** + * Builds the flow-specific token request. + * + * @return the token request to send to the authorization server + */ + protected abstract TokenRequest buildTokenRequest(); +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java new file mode 100644 index 0000000000..7e6289819c --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/ClientCredentialsTokenProvider.java @@ -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.arrow.driver.jdbc.client.oauth; + +import com.nimbusds.oauth2.sdk.ClientCredentialsGrant; +import com.nimbusds.oauth2.sdk.Scope; +import com.nimbusds.oauth2.sdk.TokenRequest; +import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic; +import com.nimbusds.oauth2.sdk.auth.Secret; +import com.nimbusds.oauth2.sdk.id.ClientID; +import java.net.URI; +import java.util.Objects; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * OAuth 2.0 Client Credentials flow token provider (RFC 6749 Section 4.4). + * + *

This provider handles service-to-service authentication where no user interaction is required. + * Tokens are cached and automatically refreshed before expiration. + */ +public class ClientCredentialsTokenProvider extends AbstractOAuthTokenProvider { + + /** + * Creates a new ClientCredentialsTokenProvider. + * + * @param tokenUri the OAuth token endpoint URI + * @param clientId the OAuth client ID + * @param clientSecret the OAuth client secret + * @param scope optional OAuth scopes (space-separated) + */ + ClientCredentialsTokenProvider( + URI tokenUri, String clientId, String clientSecret, @Nullable String scope) { + this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + Objects.requireNonNull(clientId, "clientId cannot be null"); + Objects.requireNonNull(clientSecret, "clientSecret cannot be null"); + this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret)); + this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null; + } + + @Override + protected TokenRequest buildTokenRequest() { + return new TokenRequest(tokenUri, clientAuth, new ClientCredentialsGrant(), scope); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java new file mode 100644 index 0000000000..cba9d4c2e6 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfiguration.java @@ -0,0 +1,240 @@ +/* + * 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.arrow.driver.jdbc.client.oauth; + +import com.nimbusds.oauth2.sdk.GrantType; +import java.net.URI; +import java.net.URISyntaxException; +import java.sql.SQLException; +import java.util.Locale; +import java.util.Objects; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** Configuration class for OAuth settings parsed from connection properties. */ +public class OAuthConfiguration { + + private final GrantType grantType; + private final URI tokenUri; + private final @Nullable String clientId; + private final @Nullable String clientSecret; + private final @Nullable String scope; + private final @Nullable String subjectToken; + private final @Nullable String subjectTokenType; + private final @Nullable String actorToken; + private final @Nullable String actorTokenType; + private final @Nullable String audience; + private final @Nullable String resource; + private final @Nullable String requestedTokenType; + + private OAuthConfiguration(Builder builder) throws SQLException { + this.grantType = builder.grantType; + this.tokenUri = builder.tokenUri; + this.clientId = builder.clientId; + this.clientSecret = builder.clientSecret; + this.scope = builder.scope; + this.subjectToken = builder.subjectToken; + this.subjectTokenType = builder.subjectTokenType; + this.actorToken = builder.actorToken; + this.actorTokenType = builder.actorTokenType; + this.audience = builder.audience; + this.resource = builder.resource; + this.requestedTokenType = builder.requestedTokenType; + + validate(); + } + + private void validate() throws SQLException { + Objects.requireNonNull(grantType, "OAuth grant type is required"); + Objects.requireNonNull(tokenUri, "Token URI is required"); + + if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) { + if (clientId == null || clientId.isEmpty()) { + throw new SQLException("clientId is required for client_credentials flow"); + } + if (clientSecret == null || clientSecret.isEmpty()) { + throw new SQLException("clientSecret is required for client_credentials flow"); + } + } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) { + if (subjectToken == null || subjectToken.isEmpty()) { + throw new SQLException("subjectToken is required for token_exchange flow"); + } + if (subjectTokenType == null || subjectTokenType.isEmpty()) { + throw new SQLException("subjectTokenType is required for token_exchange flow"); + } + } else { + throw new SQLException("Unsupported OAuth grant type: " + grantType); + } + } + + /** + * Creates an OAuthTokenProvider based on the configured grant type. + * + * @return the token provider + * @throws SQLException if the grant type is not supported or configuration is invalid + */ + public OAuthTokenProvider createTokenProvider() throws SQLException { + if (GrantType.CLIENT_CREDENTIALS.equals(grantType)) { + return OAuthTokenProviders.clientCredentials() + .tokenUri(tokenUri) + .clientId(clientId) + .clientSecret(clientSecret) + .scope(scope) + .build(); + } else if (GrantType.TOKEN_EXCHANGE.equals(grantType)) { + OAuthTokenProviders.TokenExchangeBuilder builder = + OAuthTokenProviders.tokenExchange() + .tokenUri(tokenUri) + .subjectToken(subjectToken) + .subjectTokenType(subjectTokenType) + .actorToken(actorToken) + .actorTokenType(actorTokenType) + .audience(audience) + .requestedTokenType(requestedTokenType) + .scope(scope) + .resource(resource); + + if (clientId != null && clientSecret != null) { + builder.clientCredentials(clientId, clientSecret); + } + + return builder.build(); + } else { + throw new SQLException("Unsupported OAuth grant type: " + grantType); + } + } + + /** Builder for OAuthConfiguration. */ + public static class Builder { + private GrantType grantType; + private URI tokenUri; + private @Nullable String clientId; + private @Nullable String clientSecret; + private @Nullable String scope; + private @Nullable String subjectToken; + private @Nullable String subjectTokenType; + private @Nullable String actorToken; + private @Nullable String actorTokenType; + private @Nullable String audience; + private @Nullable String resource; + private @Nullable String requestedTokenType; + + /** + * Sets the OAuth grant type from a string value. + * + *

Accepts either user-friendly names ("client_credentials", "token_exchange") or the full + * URN format as defined in RFC 6749 and RFC 8693. + * + * @param flowStr the flow type string (e.g., "client_credentials", "token_exchange") + * @return this builder + * @throws SQLException if the flow string is invalid + */ + public Builder flow(String flowStr) throws SQLException { + if (flowStr == null || flowStr.isEmpty()) { + throw new SQLException("OAuth flow cannot be null or empty"); + } + try { + String normalized = flowStr.toLowerCase(Locale.ROOT); + // Map user-friendly names to URN format for token_exchange + if ("token_exchange".equals(normalized)) { + normalized = GrantType.TOKEN_EXCHANGE.getValue(); + } + GrantType parsed = GrantType.parse(normalized); + if (!parsed.equals(GrantType.CLIENT_CREDENTIALS) + && !parsed.equals(GrantType.TOKEN_EXCHANGE)) { + throw new SQLException("Unsupported OAuth flow: " + flowStr); + } + this.grantType = parsed; + } catch (com.nimbusds.oauth2.sdk.ParseException e) { + throw new SQLException("Invalid OAuth flow: " + flowStr, e); + } + return this; + } + + /** + * Sets the token URI. + * + * @param tokenUri the OAuth token endpoint URI + * @return this builder + * @throws SQLException if the URI is invalid + */ + public Builder tokenUri(String tokenUri) throws SQLException { + if (tokenUri == null || tokenUri.isEmpty()) { + throw new SQLException("Token URI cannot be null or empty"); + } + try { + this.tokenUri = new URI(tokenUri); + } catch (URISyntaxException e) { + throw new SQLException("Invalid token URI: " + tokenUri, e); + } + return this; + } + + public Builder clientId(@Nullable String clientId) { + this.clientId = clientId; + return this; + } + + public Builder clientSecret(@Nullable String clientSecret) { + this.clientSecret = clientSecret; + return this; + } + + public Builder scope(@Nullable String scope) { + this.scope = scope; + return this; + } + + public Builder subjectToken(@Nullable String subjectToken) { + this.subjectToken = subjectToken; + return this; + } + + public Builder subjectTokenType(@Nullable String subjectTokenType) { + this.subjectTokenType = subjectTokenType; + return this; + } + + public Builder actorToken(@Nullable String actorToken) { + this.actorToken = actorToken; + return this; + } + + public Builder actorTokenType(@Nullable String actorTokenType) { + this.actorTokenType = actorTokenType; + return this; + } + + public Builder audience(@Nullable String audience) { + this.audience = audience; + return this; + } + + public Builder resource(@Nullable String resource) { + this.resource = resource; + return this; + } + + public Builder requestedTokenType(@Nullable String requestedTokenType) { + this.requestedTokenType = requestedTokenType; + return this; + } + + public OAuthConfiguration build() throws SQLException { + return new OAuthConfiguration(this); + } + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java new file mode 100644 index 0000000000..0d4ad4689f --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriter.java @@ -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.arrow.driver.jdbc.client.oauth; + +import java.sql.SQLException; +import java.util.Objects; +import java.util.function.Consumer; +import org.apache.arrow.flight.CallHeaders; +import org.apache.arrow.flight.auth2.Auth2Constants; + +/** Writes OAuth bearer tokens to Flight call headers. */ +public class OAuthCredentialWriter implements Consumer { + private final OAuthTokenProvider tokenProvider; + + public OAuthCredentialWriter(OAuthTokenProvider tokenProvider) { + this.tokenProvider = Objects.requireNonNull(tokenProvider, "tokenProvider cannot be null"); + } + + @Override + public void accept(CallHeaders headers) { + try { + String token = tokenProvider.getValidToken(); + headers.insert(Auth2Constants.AUTHORIZATION_HEADER, Auth2Constants.BEARER_PREFIX + token); + } catch (SQLException e) { + throw new OAuthTokenException("Failed to obtain OAuth token", e); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java similarity index 64% rename from vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java rename to flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java index 1b1bf4e6e4..aceadb327b 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java @@ -14,18 +14,18 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.arrow.vector.complex.impl; +package org.apache.arrow.driver.jdbc.client.oauth; -import org.apache.arrow.vector.ExtensionTypeVector; -import org.apache.arrow.vector.UuidVector; - -public class UuidWriterFactory implements ExtensionTypeWriterFactory { +/** + * Runtime exception thrown when OAuth token operations fail. Used to wrap checked exceptions in + * contexts that don't allow them. + */ +public class OAuthTokenException extends RuntimeException { + public OAuthTokenException(String message) { + super(message); + } - @Override - public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) { - if (extensionTypeVector instanceof UuidVector) { - return new UuidWriterImpl((UuidVector) extensionTypeVector); - } - return null; + public OAuthTokenException(String message, Throwable cause) { + super(message, cause); } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java new file mode 100644 index 0000000000..241611e432 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProvider.java @@ -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.arrow.driver.jdbc.client.oauth; + +import java.sql.SQLException; + +/** + * Interface for OAuth token providers that handle token acquisition and refresh. Implementations + * should cache tokens and automatically refresh them before expiration. + */ +public interface OAuthTokenProvider { + /** + * Gets a valid OAuth access token, refreshing if necessary. + * + * @return a valid access token string + * @throws SQLException if token cannot be obtained + */ + String getValidToken() throws SQLException; +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java new file mode 100644 index 0000000000..bbf7072d39 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenProviders.java @@ -0,0 +1,419 @@ +/* + * 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.arrow.driver.jdbc.client.oauth; + +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.Scope; +import com.nimbusds.oauth2.sdk.auth.ClientAuthentication; +import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic; +import com.nimbusds.oauth2.sdk.auth.Secret; +import com.nimbusds.oauth2.sdk.id.Audience; +import com.nimbusds.oauth2.sdk.id.ClientID; +import com.nimbusds.oauth2.sdk.token.TokenTypeURI; +import com.nimbusds.oauth2.sdk.token.TypelessAccessToken; +import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant; +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Unified factory for creating OAuth token providers. + * + *

This class provides a single entry point for creating all OAuth token providers with a + * consistent builder API. It supports: + * + *

    + *
  • Client Credentials flow (RFC 6749 Section 4.4) + *
  • Token Exchange flow (RFC 8693) + *
+ * + *

Example usage: + * + * 

{@code
+ * // Client Credentials flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.clientCredentials()
+ *     .tokenUri("https://auth.example.com/token")
+ *     .clientId("my-client")
+ *     .clientSecret("my-secret")
+ *     .scope("read write")
+ *     .build();
+ *
+ * // Token Exchange flow
+ * OAuthTokenProvider provider = OAuthTokenProviders.tokenExchange()
+ *     .tokenUri("https://auth.example.com/token")
+ *     .subjectToken("user-token")
+ *     .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")
+ *     .build();
+ * }
+ */ +public final class OAuthTokenProviders { + + private OAuthTokenProviders() {} + + /** + * Creates a new builder for Client Credentials flow. + * + * @return a new ClientCredentialsBuilder instance + */ + public static ClientCredentialsBuilder clientCredentials() { + return new ClientCredentialsBuilder(); + } + + /** + * Creates a new builder for Token Exchange flow. + * + * @return a new TokenExchangeBuilder instance + */ + public static TokenExchangeBuilder tokenExchange() { + return new TokenExchangeBuilder(); + } + + /** Builder for creating {@link ClientCredentialsTokenProvider} instances. */ + public static class ClientCredentialsBuilder { + private @Nullable URI tokenUri; + private @Nullable String clientId; + private @Nullable String clientSecret; + private @Nullable String scope; + + ClientCredentialsBuilder() {} + + /** + * Sets the OAuth token endpoint URI (required). + * + * @param tokenUri the token endpoint URI + * @return this builder + */ + public ClientCredentialsBuilder tokenUri(URI tokenUri) { + this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + return this; + } + + /** + * Sets the OAuth token endpoint URI from a string (required). + * + * @param tokenUri the token endpoint URI string + * @return this builder + * @throws IllegalArgumentException if the URI is invalid + */ + public ClientCredentialsBuilder tokenUri(String tokenUri) { + Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + try { + this.tokenUri = new URI(tokenUri); + } catch (URISyntaxException e) { + throw new IllegalArgumentException("Invalid token URI: " + tokenUri, e); + } + return this; + } + + /** + * Sets the OAuth client ID (required). + * + * @param clientId the client ID + * @return this builder + */ + public ClientCredentialsBuilder clientId(String clientId) { + this.clientId = Objects.requireNonNull(clientId, "clientId cannot be null"); + return this; + } + + /** + * Sets the OAuth client secret (required). + * + * @param clientSecret the client secret + * @return this builder + */ + public ClientCredentialsBuilder clientSecret(String clientSecret) { + this.clientSecret = Objects.requireNonNull(clientSecret, "clientSecret cannot be null"); + return this; + } + + /** + * Sets the OAuth scopes (optional). + * + * @param scope the space-separated scope string + * @return this builder + */ + public ClientCredentialsBuilder scope(@Nullable String scope) { + this.scope = scope; + return this; + } + + /** + * Builds a new ClientCredentialsTokenProvider instance. + * + * @return the configured ClientCredentialsTokenProvider + * @throws IllegalStateException if required parameters are missing + */ + public ClientCredentialsTokenProvider build() { + if (tokenUri == null) { + throw new IllegalStateException("tokenUri is required"); + } + if (clientId == null) { + throw new IllegalStateException("clientId is required"); + } + if (clientSecret == null) { + throw new IllegalStateException("clientSecret is required"); + } + return new ClientCredentialsTokenProvider(tokenUri, clientId, clientSecret, scope); + } + } + + /** Builder for creating {@link TokenExchangeTokenProvider} instances. */ + public static class TokenExchangeBuilder { + private @Nullable URI tokenUri; + private @Nullable String subjectToken; + private @Nullable String subjectTokenType; + private @Nullable String actorToken; + private @Nullable String actorTokenType; + private @Nullable String audience; + private @Nullable String requestedTokenType; + private @Nullable Scope scope; + private @Nullable List resources; + private @Nullable ClientAuthentication clientAuth; + + TokenExchangeBuilder() {} + + /** + * Sets the OAuth token endpoint URI (required). + * + * @param tokenUri the token endpoint URI + * @return this builder + */ + public TokenExchangeBuilder tokenUri(URI tokenUri) { + this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + return this; + } + + /** + * Sets the OAuth token endpoint URI from a string (required). + * + * @param tokenUri the token endpoint URI string + * @return this builder + * @throws IllegalArgumentException if the URI is invalid + */ + public TokenExchangeBuilder tokenUri(String tokenUri) { + Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + try { + this.tokenUri = new URI(tokenUri); + } catch (URISyntaxException e) { + throw new IllegalArgumentException("Invalid token URI: " + tokenUri, e); + } + return this; + } + + /** + * Sets the subject token to exchange (required). + * + * @param subjectToken the subject token value + * @return this builder + */ + public TokenExchangeBuilder subjectToken(String subjectToken) { + this.subjectToken = Objects.requireNonNull(subjectToken, "subjectToken cannot be null"); + return this; + } + + /** + * Sets the type of the subject token (required). + * + * @param subjectTokenType the subject token type URI + * @return this builder + */ + public TokenExchangeBuilder subjectTokenType(String subjectTokenType) { + this.subjectTokenType = + Objects.requireNonNull(subjectTokenType, "subjectTokenType cannot be null"); + return this; + } + + /** + * Sets the optional actor token for delegation scenarios. + * + * @param actorToken the actor token value + * @return this builder + */ + public TokenExchangeBuilder actorToken(@Nullable String actorToken) { + this.actorToken = actorToken; + return this; + } + + /** + * Sets the type of the actor token. + * + * @param actorTokenType the actor token type URI + * @return this builder + */ + public TokenExchangeBuilder actorTokenType(@Nullable String actorTokenType) { + this.actorTokenType = actorTokenType; + return this; + } + + /** + * Sets the target audience for the exchanged token. + * + * @param audience the target audience + * @return this builder + */ + public TokenExchangeBuilder audience(@Nullable String audience) { + this.audience = audience; + return this; + } + + /** + * Sets the requested token type for the exchanged token. + * + * @param requestedTokenType the requested token type URI + * @return this builder + */ + public TokenExchangeBuilder requestedTokenType(@Nullable String requestedTokenType) { + this.requestedTokenType = requestedTokenType; + return this; + } + + /** + * Sets the OAuth scopes for the token request. + * + * @param scope the OAuth scope object + * @return this builder + */ + public TokenExchangeBuilder scope(@Nullable Scope scope) { + this.scope = scope; + return this; + } + + /** + * Sets the OAuth scopes from a space-separated string. + * + * @param scope the space-separated scope string + * @return this builder + */ + public TokenExchangeBuilder scope(@Nullable String scope) { + this.scope = (scope != null && !scope.isEmpty()) ? Scope.parse(scope) : null; + return this; + } + + /** + * Sets the target resource URIs (RFC 8707). + * + * @param resources the list of resource URIs + * @return this builder + */ + public TokenExchangeBuilder resources(@Nullable List resources) { + this.resources = resources; + return this; + } + + /** + * Sets a single target resource URI (RFC 8707). + * + * @param resource the resource URI + * @return this builder + */ + public TokenExchangeBuilder resource(@Nullable URI resource) { + this.resources = resource != null ? Collections.singletonList(resource) : null; + return this; + } + + /** + * Sets a single target resource URI from a string (RFC 8707). + * + * @param resource the resource URI string + * @return this builder + */ + public TokenExchangeBuilder resource(@Nullable String resource) { + if (resource != null && !resource.isEmpty()) { + this.resources = Collections.singletonList(URI.create(resource)); + } else { + this.resources = null; + } + return this; + } + + /** + * Sets the client authentication. + * + * @param clientAuth the client authentication object + * @return this builder + */ + public TokenExchangeBuilder clientAuthentication(@Nullable ClientAuthentication clientAuth) { + this.clientAuth = clientAuth; + return this; + } + + /** + * Sets client authentication using client ID and secret. + * + * @param clientId the client ID + * @param clientSecret the client secret + * @return this builder + */ + public TokenExchangeBuilder clientCredentials(String clientId, String clientSecret) { + Objects.requireNonNull(clientId, "clientId cannot be null"); + Objects.requireNonNull(clientSecret, "clientSecret cannot be null"); + this.clientAuth = new ClientSecretBasic(new ClientID(clientId), new Secret(clientSecret)); + return this; + } + + /** + * Builds a new TokenExchangeTokenProvider instance. + * + * @return the configured TokenExchangeTokenProvider + * @throws IllegalStateException if required parameters are missing + */ + public TokenExchangeTokenProvider build() { + if (tokenUri == null) { + throw new IllegalStateException("tokenUri is required"); + } + if (subjectToken == null) { + throw new IllegalStateException("subjectToken is required"); + } + if (subjectTokenType == null) { + throw new IllegalStateException("subjectTokenType is required"); + } + + TokenExchangeGrant grant = createGrant(); + return new TokenExchangeTokenProvider(tokenUri, grant, clientAuth, scope, resources); + } + + private TokenExchangeGrant createGrant() { + try { + TypelessAccessToken subjectAccessToken = new TypelessAccessToken(subjectToken); + TokenTypeURI subjectTypeUri = TokenTypeURI.parse(subjectTokenType); + + TypelessAccessToken actorAccessToken = + actorToken != null ? new TypelessAccessToken(actorToken) : null; + TokenTypeURI actorTypeUri = + actorTokenType != null ? TokenTypeURI.parse(actorTokenType) : null; + TokenTypeURI requestedTypeUri = + requestedTokenType != null ? TokenTypeURI.parse(requestedTokenType) : null; + List audienceList = + audience != null ? Collections.singletonList(new Audience(audience)) : null; + + return new TokenExchangeGrant( + subjectAccessToken, + subjectTypeUri, + actorAccessToken, + actorTypeUri, + requestedTypeUri, + audienceList); + } catch (ParseException e) { + throw new IllegalStateException("Failed to create TokenExchangeGrant", e); + } + } + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java new file mode 100644 index 0000000000..af433a2712 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenExchangeTokenProvider.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.client.oauth; + +import com.nimbusds.oauth2.sdk.Scope; +import com.nimbusds.oauth2.sdk.TokenRequest; +import com.nimbusds.oauth2.sdk.auth.ClientAuthentication; +import com.nimbusds.oauth2.sdk.tokenexchange.TokenExchangeGrant; +import java.net.URI; +import java.util.List; +import java.util.Objects; +import org.apache.arrow.util.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * OAuth 2.0 Token Exchange flow token provider (RFC 8693). + * + *

This provider exchanges one token for another, commonly used for federated authentication, + * delegation, or impersonation scenarios. Tokens are cached and automatically refreshed. + */ +public class TokenExchangeTokenProvider extends AbstractOAuthTokenProvider { + + @VisibleForTesting TokenExchangeGrant grant; + + @VisibleForTesting @Nullable List resources; + + /** + * Creates a new TokenExchangeTokenProvider with full configuration. + * + * @param tokenUri the OAuth token endpoint URI + * @param grant the token exchange grant containing subject/actor token information + * @param clientAuth optional client authentication + * @param scope optional OAuth scopes + * @param resource optional target resource URI (RFC 8707) + */ + TokenExchangeTokenProvider( + URI tokenUri, + TokenExchangeGrant grant, + @Nullable ClientAuthentication clientAuth, + @Nullable Scope scope, + @Nullable List resource) { + this.tokenUri = Objects.requireNonNull(tokenUri, "tokenUri cannot be null"); + this.grant = Objects.requireNonNull(grant, "grant cannot be null"); + this.scope = scope; + this.resources = resource; + this.clientAuth = clientAuth; + } + + @Override + protected TokenRequest buildTokenRequest() { + TokenRequest.Builder builder; + if (clientAuth != null) { + builder = new TokenRequest.Builder(tokenUri, clientAuth, grant); + } else { + builder = new TokenRequest.Builder(tokenUri, grant); + } + + if (scope != null) { + builder.scope(scope); + } + if (resources != null) { + builder.resources(resources.toArray(new URI[0])); + } + + return builder.build(); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java new file mode 100644 index 0000000000..f47cc8b053 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/TokenInfo.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.client.oauth; + +import java.time.Instant; +import java.util.Objects; + +/** Holds OAuth token information including the access token and expiration time. */ +public class TokenInfo { + private final String accessToken; + private final Instant expiresAt; + + public TokenInfo(String accessToken, Instant expiresAt) { + this.accessToken = Objects.requireNonNull(accessToken, "accessToken cannot be null"); + this.expiresAt = Objects.requireNonNull(expiresAt, "expiresAt cannot be null"); + } + + public String getAccessToken() { + return accessToken; + } + + /** + * Checks if the token is expired or will expire within the buffer period. + * + * @param bufferSeconds seconds before actual expiration to consider token expired + * @return true if token should be refreshed + */ + public boolean isExpired(int bufferSeconds) { + return Instant.now().plusSeconds(bufferSeconds).isAfter(expiresAt); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCache.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCache.java new file mode 100644 index 0000000000..36e8441baa --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCache.java @@ -0,0 +1,55 @@ +/* + * 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.arrow.driver.jdbc.client.utils; + +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; +import java.time.Duration; +import org.apache.arrow.util.VisibleForTesting; + +/** + * A cache for Flight clients. + * + *

The intent is to avoid constantly recreating clients to the same locations. gRPC can multiplex + * multiple requests over a single TCP connection, and a cache would let us take advantage of that. + * + *

At the time being it only tracks whether a location is reachable or not. To actually cache + * clients, we would need a way to incorporate other connection parameters (authentication, etc.) + * into the cache key. + */ +public final class FlightClientCache { + @VisibleForTesting Cache clientCache; + + public FlightClientCache() { + this.clientCache = Caffeine.newBuilder().expireAfterWrite(Duration.ofSeconds(600)).build(); + } + + public boolean isDud(String key) { + return clientCache.getIfPresent(key) != null; + } + + public void markLocationAsDud(String key) { + clientCache.put(key, new ClientCacheEntry()); + } + + public void markLocationAsReachable(String key) { + clientCache.invalidate(key); + } + + /** A cache entry (empty because we only track reachability, see outer class docstring). */ + public static final class ClientCacheEntry {} +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueue.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueue.java new file mode 100644 index 0000000000..f507ec53e7 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueue.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.client.utils; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.flight.Location; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A queue of Flight locations to connect to for an endpoint. + * + *

This helper class is intended to encapsulate the retry logic in a testable manner. + */ +public final class FlightLocationQueue implements Iterator { + private final Deque locations; + private final Deque badLocations; + + /** + * Create a new queue. + * + * @param flightClientCache An optional cache used to sort previously unreachable locations to the + * end. + * @param locations The locations to try. + */ + public FlightLocationQueue( + @Nullable FlightClientCache flightClientCache, List locations) { + this.locations = new ArrayDeque<>(); + this.badLocations = new ArrayDeque<>(); + + for (Location location : locations) { + if (flightClientCache != null && flightClientCache.isDud(location.toString())) { + this.badLocations.add(location); + } else { + this.locations.add(location); + } + } + } + + @Override + public boolean hasNext() { + return !locations.isEmpty() || !badLocations.isEmpty(); + } + + @Override + public Location next() { + if (!locations.isEmpty()) { + return locations.pop(); + } else if (!badLocations.isEmpty()) { + return badLocations.pop(); + } + throw new NoSuchElementException(); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java index a035bbba49..d692f39372 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/BinaryViewAvaticaParameterConverter.java @@ -17,6 +17,7 @@ package org.apache.arrow.driver.jdbc.converter.impl; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ViewVarBinaryVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; @@ -29,7 +30,12 @@ public BinaryViewAvaticaParameterConverter(ArrowType.BinaryView type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { - throw new UnsupportedOperationException("Not implemented"); + byte[] value = (byte[]) typedValue.toJdbc(null); + if (vector instanceof ViewVarBinaryVector) { + ((ViewVarBinaryVector) vector).setSafe(index, value); + return true; + } + return false; } @Override diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java index 076fefc42a..c9d9f2926b 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/Utf8ViewAvaticaParameterConverter.java @@ -17,8 +17,10 @@ package org.apache.arrow.driver.jdbc.converter.impl; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.Text; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.remote.TypedValue; @@ -29,7 +31,12 @@ public Utf8ViewAvaticaParameterConverter(ArrowType.Utf8View type) {} @Override public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { - throw new UnsupportedOperationException("Utf8View not supported"); + String value = (String) typedValue.toLocal(); + if (vector instanceof ViewVarCharVector) { + ((ViewVarCharVector) vector).setSafe(index, new Text(value)); + return true; + } + return false; } @Override diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java new file mode 100644 index 0000000000..b2157890cf --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverter.java @@ -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.arrow.driver.jdbc.converter.impl; + +import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeIdFromArrowType; +import static org.apache.arrow.driver.jdbc.utils.SqlTypes.getSqlTypeNameFromArrowType; + +import java.nio.ByteBuffer; +import java.util.UUID; +import org.apache.arrow.driver.jdbc.converter.AvaticaParameterConverter; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.util.UuidUtility; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.remote.TypedValue; +import org.apache.calcite.avatica.util.ByteString; + +/** + * AvaticaParameterConverter for UUID Arrow extension type. + * + *

Handles conversion of UUID values from JDBC parameters to Arrow's UUID extension type. Accepts + * both {@link UUID} objects and String representations of UUIDs. + */ +public class UuidAvaticaParameterConverter implements AvaticaParameterConverter { + + public UuidAvaticaParameterConverter() {} + + @Override + public boolean bindParameter(FieldVector vector, TypedValue typedValue, int index) { + if (!(vector instanceof UuidVector)) { + return false; + } + + UuidVector uuidVector = (UuidVector) vector; + Object value = typedValue.toJdbc(null); + + if (value == null) { + uuidVector.setNull(index); + return true; + } + + UUID uuid; + if (value instanceof UUID) { + uuid = (UUID) value; + } else if (value instanceof String) { + uuid = UUID.fromString((String) value); + } else if (value instanceof byte[]) { + byte[] bytes = (byte[]) value; + if (bytes.length != 16) { + throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length); + } + uuid = uuidFromBytes(bytes); + } else if (value instanceof ByteString) { + byte[] bytes = ((ByteString) value).getBytes(); + if (bytes.length != 16) { + throw new IllegalArgumentException("UUID byte array must be 16 bytes, got " + bytes.length); + } + uuid = uuidFromBytes(bytes); + } else { + throw new IllegalArgumentException( + "Cannot convert " + value.getClass().getName() + " to UUID"); + } + + uuidVector.setSafe(index, UuidUtility.getBytesFromUUID(uuid)); + return true; + } + + @Override + public AvaticaParameter createParameter(Field field) { + final String name = field.getName(); + final int jdbcType = getSqlTypeIdFromArrowType(field.getType()); + final String typeName = getSqlTypeNameFromArrowType(field.getType()); + final String className = UUID.class.getCanonicalName(); + return new AvaticaParameter(false, 0, 0, jdbcType, typeName, className, name); + } + + private static UUID uuidFromBytes(byte[] bytes) { + final long mostSignificantBits; + final long leastSignificantBits; + ByteBuffer bb = ByteBuffer.wrap(bytes); + // Reads the first eight bytes + mostSignificantBits = bb.getLong(); + // Reads the first eight bytes at this buffer's current + leastSignificantBits = bb.getLong(); + + return new UUID(mostSignificantBits, leastSignificantBits); + } +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java index ab6a5898b7..d0ba74dbcc 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.driver.jdbc.utils; +import java.sql.SQLException; import java.time.Duration; import java.util.Arrays; import java.util.HashMap; @@ -23,6 +24,7 @@ import java.util.Objects; import java.util.Properties; import org.apache.arrow.driver.jdbc.ArrowFlightConnection; +import org.apache.arrow.driver.jdbc.client.oauth.OAuthConfiguration; import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.FlightCallHeaders; @@ -31,6 +33,7 @@ import org.apache.calcite.avatica.ConnectionConfig; import org.apache.calcite.avatica.ConnectionConfigImpl; import org.apache.calcite.avatica.ConnectionProperty; +import org.checkerframework.checker.nullness.qual.Nullable; /** A {@link ConnectionConfig} for the {@link ArrowFlightConnection}. */ public final class ArrowFlightConnectionConfigImpl extends ConnectionConfigImpl { @@ -174,6 +177,11 @@ public Duration getConnectTimeout() { return Duration.ofMillis(timeout); } + /** Whether to enable the client cache. */ + public boolean useClientCache() { + return ArrowFlightConnectionProperty.USE_CLIENT_CACHE.getBoolean(properties); + } + /** * Gets the {@link CallOption}s from this {@link ConnectionConfig}. * @@ -206,6 +214,38 @@ public Map getHeaderAttributes() { return headers; } + /** + * Returns OAuth configuration if oauth.flow is specified, null otherwise. + * + * @return the OAuth configuration or null + * @throws SQLException if the OAuth configuration is invalid + */ + public @Nullable OAuthConfiguration getOauthConfiguration() throws SQLException { + String flow = ArrowFlightConnectionProperty.OAUTH_FLOW.getString(properties); + if (flow == null) { + return null; + } + + return new OAuthConfiguration.Builder() + .flow(flow) + .clientId(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.getString(properties)) + .clientSecret(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.getString(properties)) + .tokenUri(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.getString(properties)) + .scope(ArrowFlightConnectionProperty.OAUTH_SCOPE.getString(properties)) + .resource(ArrowFlightConnectionProperty.OAUTH_RESOURCE.getString(properties)) + .subjectToken( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.getString(properties)) + .subjectTokenType( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.getString(properties)) + .actorToken(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.getString(properties)) + .actorTokenType( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.getString(properties)) + .audience(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.getString(properties)) + .requestedTokenType( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.getString(properties)) + .build(); + } + /** Custom {@link ConnectionProperty} for the {@link ArrowFlightConnectionConfigImpl}. */ public enum ArrowFlightConnectionProperty implements ConnectionProperty { HOST("host", null, Type.STRING, true), @@ -226,6 +266,24 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty { RETAIN_AUTH("retainAuth", true, Type.BOOLEAN, false), CATALOG("catalog", null, Type.STRING, false), CONNECT_TIMEOUT_MILLIS("connectTimeoutMs", 10000, Type.NUMBER, false), + USE_CLIENT_CACHE("useClientCache", true, Type.BOOLEAN, false), + + // OAuth configuration properties + OAUTH_FLOW("oauth.flow", null, Type.STRING, false), + OAUTH_CLIENT_ID("oauth.clientId", null, Type.STRING, false), + OAUTH_CLIENT_SECRET("oauth.clientSecret", null, Type.STRING, false), + OAUTH_TOKEN_URI("oauth.tokenUri", null, Type.STRING, false), + OAUTH_SCOPE("oauth.scope", null, Type.STRING, false), + OAUTH_RESOURCE("oauth.resource", null, Type.STRING, false), + + // Token exchange specific properties + OAUTH_EXCHANGE_SUBJECT_TOKEN("oauth.exchange.subjectToken", null, Type.STRING, false), + OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE("oauth.exchange.subjectTokenType", null, Type.STRING, false), + OAUTH_EXCHANGE_ACTOR_TOKEN("oauth.exchange.actorToken", null, Type.STRING, false), + OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE("oauth.exchange.actorTokenType", null, Type.STRING, false), + OAUTH_EXCHANGE_AUDIENCE("oauth.exchange.aud", null, Type.STRING, false), + OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE( + "oauth.exchange.requestedTokenType", null, Type.STRING, false), ; private final String camelName; diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java index 4c2a9b865f..8f40d6698e 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java @@ -19,6 +19,7 @@ import java.util.List; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; import org.apache.arrow.driver.jdbc.converter.impl.BinaryAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.BinaryViewAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.BoolAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.DateAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.DecimalAvaticaParameterConverter; @@ -39,11 +40,17 @@ import org.apache.arrow.driver.jdbc.converter.impl.TimestampAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.calcite.avatica.remote.TypedValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Convert Avatica PreparedStatement parameters from a list of TypedValue to Arrow and bind them to @@ -108,9 +115,9 @@ public void bind(List typedValues, int index) { * @param typedValue TypedValue to bind to the vector. * @param index Vector index to bind the value at. */ - private void bind(FieldVector vector, TypedValue typedValue, int index) { + private void bind(FieldVector vector, @Nullable TypedValue typedValue, int index) { try { - if (typedValue.value == null) { + if (typedValue == null || typedValue.value == null) { if (vector.getField().isNullable()) { vector.setNull(index); } else { @@ -127,7 +134,7 @@ private void bind(FieldVector vector, TypedValue typedValue, int index) { throw new UnsupportedOperationException( String.format( "Binding value of type %s is not yet supported for expected Arrow type %s", - typedValue.type, vector.getField().getType())); + typedValue == null ? "null" : typedValue.type, vector.getField().getType())); } } @@ -207,7 +214,7 @@ public Boolean visit(ArrowType.Utf8 type) { @Override public Boolean visit(ArrowType.Utf8View type) { - throw new UnsupportedOperationException("Utf8View is unsupported"); + return new Utf8ViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override @@ -222,7 +229,7 @@ public Boolean visit(ArrowType.Binary type) { @Override public Boolean visit(ArrowType.BinaryView type) { - throw new UnsupportedOperationException("BinaryView is unsupported"); + return new BinaryViewAvaticaParameterConverter(type).bindParameter(vector, typedValue, index); } @Override @@ -287,5 +294,15 @@ public Boolean visit(ArrowType.RunEndEncoded type) { throw new UnsupportedOperationException( "No Avatica parameter binder implemented for type " + type); } + + @Override + public Boolean visit(ExtensionType type) { + if (type instanceof UuidType) { + return new UuidAvaticaParameterConverter().bindParameter(vector, typedValue, index); + } + + // fallback to default implementation + return ArrowTypeVisitor.super.visit(type); + } } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java index 17b0f42dc7..dd51ee5361 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ConvertUtils.java @@ -43,8 +43,12 @@ import org.apache.arrow.driver.jdbc.converter.impl.UnionAvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8AvaticaParameterConverter; import org.apache.arrow.driver.jdbc.converter.impl.Utf8ViewAvaticaParameterConverter; +import org.apache.arrow.driver.jdbc.converter.impl.UuidAvaticaParameterConverter; import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeVisitor; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.calcite.avatica.AvaticaParameter; import org.apache.calcite.avatica.ColumnMetaData; @@ -136,6 +140,10 @@ public static void setOnColumnMetaDataBuilder( if (searchable != null) { builder.setSearchable(searchable); } + final String remarks = columnMetadata.getRemarks(); + if (remarks != null) { + builder.setLabel(remarks); + } } /** @@ -290,5 +298,15 @@ public AvaticaParameter visit(ArrowType.RunEndEncoded type) { throw new UnsupportedOperationException( "No Avatica parameter binder implemented for type " + type); } + + @Override + public AvaticaParameter visit(ExtensionType type) { + if (type instanceof UuidType) { + return new UuidAvaticaParameterConverter().createParameter(field); + } + + // fallback to default implementation + return ArrowTypeVisitor.super.visit(type); + } } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java index 1b76ca0c95..7982d5bc73 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/SqlTypes.java @@ -20,11 +20,13 @@ import java.sql.Types; import java.util.HashMap; import java.util.Map; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.pojo.ArrowType; /** SQL Types utility functions. */ public class SqlTypes { + private static final Map typeIdToName = new HashMap<>(); static { @@ -107,12 +109,17 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) { } break; case Binary: + case BinaryView: return Types.VARBINARY; case FixedSizeBinary: + if (arrowType instanceof UuidType) { + return Types.OTHER; + } return Types.BINARY; case LargeBinary: return Types.LONGVARBINARY; case Utf8: + case Utf8View: return Types.VARCHAR; case LargeUtf8: return Types.LONGVARCHAR; diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java index 70d3bcbd33..3ab1460b27 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowDatabaseMetadataTest.java @@ -26,7 +26,6 @@ import static java.util.stream.Collectors.toList; import static java.util.stream.IntStream.range; import static org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer.serializeSchema; -import static org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; import static org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportsConvert.SQL_CONVERT_BIGINT_VALUE; import static org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportsConvert.SQL_CONVERT_BIT_VALUE; import static org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportsConvert.SQL_CONVERT_INTEGER_VALUE; @@ -55,9 +54,11 @@ import org.apache.arrow.driver.jdbc.utils.ResultSetTestUtils; import org.apache.arrow.driver.jdbc.utils.ThrowableAssertionUtils; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; +import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; import org.apache.arrow.flight.sql.impl.FlightSql; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; @@ -79,6 +80,7 @@ import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.AfterAll; @@ -322,7 +324,7 @@ public class ArrowDatabaseMetadataTest { expectedGetColumnsDecimalDigits.get(i % 3), expectedGetColumnsRadix.get(i % 3), !Objects.equals(expectedGetColumnsIsNullable.get(i % 3), "NO") ? 1 : 0, - null, + format("column description #%d", (i % 3) + 1), null, null, null, @@ -419,17 +421,44 @@ public static void setUpBeforeClass() throws SQLException { try (final BufferAllocator allocator = new RootAllocator(); final VectorSchemaRoot root = VectorSchemaRoot.create(Schemas.GET_TABLES_SCHEMA, allocator)) { + final Field field1 = + new Field( + "column_1", + new FieldType( + true, + ArrowType.Decimal.createDecimal(5, 2, 128), + null, + new FlightSqlColumnMetadata.Builder() + .remarks("column description #1") + .build() + .getMetadataMap()), + null); + final Field field2 = + new Field( + "column_2", + new FieldType( + true, + new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), + null, + new FlightSqlColumnMetadata.Builder() + .remarks("column description #2") + .build() + .getMetadataMap()), + null); + final Field field3 = + new Field( + "column_3", + new FieldType( + false, + Types.MinorType.INT.getType(), + null, + new FlightSqlColumnMetadata.Builder() + .remarks("column description #3") + .build() + .getMetadataMap()), + null); final byte[] filledTableSchemaBytes = - copyFrom( - serializeSchema( - new Schema( - Arrays.asList( - Field.nullable( - "column_1", ArrowType.Decimal.createDecimal(5, 2, 128)), - Field.nullable( - "column_2", - new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), - Field.notNullable("column_3", Types.MinorType.INT.getType()))))) + copyFrom(serializeSchema(new Schema(Arrays.asList(field1, field2, field3)))) .toByteArray(); final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); @@ -1514,11 +1543,83 @@ public void testEmptySqlInfo() throws Exception { try (final Connection testConnection = FLIGHT_SERVER_EMPTY_SQLINFO_TEST_RULE.getConnection(false)) { final DatabaseMetaData metaData = testConnection.getMetaData(); + assertThat(metaData.getSQLKeywords(), is("")); assertThat(metaData.getNumericFunctions(), is("")); assertThat(metaData.getStringFunctions(), is("")); assertThat(metaData.getSystemFunctions(), is("")); assertThat(metaData.getTimeDateFunctions(), is("")); + + assertThat(metaData.getMaxBinaryLiteralLength(), is(0)); + assertThat(metaData.getMaxCharLiteralLength(), is(0)); + assertThat(metaData.getMaxColumnNameLength(), is(0)); + assertThat(metaData.getMaxColumnsInGroupBy(), is(0)); + assertThat(metaData.getMaxColumnsInIndex(), is(0)); + assertThat(metaData.getMaxColumnsInOrderBy(), is(0)); + assertThat(metaData.getMaxColumnsInSelect(), is(0)); + assertThat(metaData.getMaxColumnsInTable(), is(0)); + assertThat(metaData.getMaxConnections(), is(0)); + assertThat(metaData.getMaxCursorNameLength(), is(0)); + assertThat(metaData.getMaxIndexLength(), is(0)); + assertThat(metaData.getMaxSchemaNameLength(), is(0)); + assertThat(metaData.getMaxProcedureNameLength(), is(0)); + assertThat(metaData.getMaxCatalogNameLength(), is(0)); + assertThat(metaData.getMaxRowSize(), is(0)); + assertThat(metaData.getMaxStatementLength(), is(0)); + assertThat(metaData.getMaxStatements(), is(0)); + assertThat(metaData.getMaxTableNameLength(), is(0)); + assertThat(metaData.getMaxTablesInSelect(), is(0)); + assertThat(metaData.getMaxUserNameLength(), is(0)); + + assertThat(metaData.supportsColumnAliasing(), is(false)); + assertThat(metaData.nullPlusNonNullIsNull(), is(false)); + assertThat(metaData.supportsTableCorrelationNames(), is(false)); + assertThat(metaData.supportsDifferentTableCorrelationNames(), is(false)); + assertThat(metaData.supportsExpressionsInOrderBy(), is(false)); + assertThat(metaData.supportsOrderByUnrelated(), is(false)); + assertThat(metaData.supportsLikeEscapeClause(), is(false)); + assertThat(metaData.supportsNonNullableColumns(), is(false)); + assertThat(metaData.supportsIntegrityEnhancementFacility(), is(false)); + assertThat(metaData.isCatalogAtStart(), is(false)); + assertThat(metaData.supportsSelectForUpdate(), is(false)); + assertThat(metaData.supportsStoredProcedures(), is(false)); + assertThat(metaData.supportsCorrelatedSubqueries(), is(false)); + assertThat(metaData.doesMaxRowSizeIncludeBlobs(), is(false)); + assertThat(metaData.supportsTransactions(), is(false)); + assertThat(metaData.dataDefinitionCausesTransactionCommit(), is(false)); + assertThat(metaData.dataDefinitionIgnoredInTransactions(), is(false)); + assertThat(metaData.supportsBatchUpdates(), is(false)); + assertThat(metaData.supportsSavepoints(), is(false)); + assertThat(metaData.supportsNamedParameters(), is(false)); + assertThat(metaData.locatorsUpdateCopy(), is(false)); + assertThat(metaData.supportsStoredFunctionsUsingCallSyntax(), is(false)); + assertThat(metaData.supportsGroupBy(), is(false)); + assertThat(metaData.supportsGroupByUnrelated(), is(false)); + assertThat(metaData.supportsMinimumSQLGrammar(), is(false)); + assertThat(metaData.supportsCoreSQLGrammar(), is(false)); + assertThat(metaData.supportsExtendedSQLGrammar(), is(false)); + assertThat(metaData.supportsANSI92EntryLevelSQL(), is(false)); + assertThat(metaData.supportsANSI92IntermediateSQL(), is(false)); + assertThat(metaData.supportsANSI92FullSQL(), is(false)); + assertThat(metaData.supportsOuterJoins(), is(false)); + assertThat(metaData.supportsFullOuterJoins(), is(false)); + assertThat(metaData.supportsLimitedOuterJoins(), is(false)); + assertThat(metaData.supportsSchemasInProcedureCalls(), is(false)); + assertThat(metaData.supportsSchemasInIndexDefinitions(), is(false)); + assertThat(metaData.supportsSchemasInPrivilegeDefinitions(), is(false)); + assertThat(metaData.supportsCatalogsInIndexDefinitions(), is(false)); + assertThat(metaData.supportsCatalogsInPrivilegeDefinitions(), is(false)); + assertThat(metaData.supportsPositionedDelete(), is(false)); + assertThat(metaData.supportsPositionedUpdate(), is(false)); + assertThat(metaData.supportsSubqueriesInComparisons(), is(false)); + assertThat(metaData.supportsSubqueriesInExists(), is(false)); + assertThat(metaData.supportsSubqueriesInIns(), is(false)); + assertThat(metaData.supportsSubqueriesInQuantifieds(), is(false)); + assertThat(metaData.supportsUnion(), is(false)); + assertThat(metaData.supportsUnionAll(), is(false)); + assertThat(metaData.supportsConvert(), is(false)); + + assertThat(metaData.getDefaultTransactionIsolation(), is(Connection.TRANSACTION_NONE)); } } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java index 06d101724c..cb6abacb2f 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcArrayTest.java @@ -129,7 +129,7 @@ public void testShouldGetResultSetReturnValidResultSet() throws SQLException { try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet()) { int count = 0; while (resultSet.next()) { - assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count)); + assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count)); count++; } } @@ -142,7 +142,7 @@ public void testShouldGetResultSetReturnValidResultSetWithOffsets() throws SQLEx try (ResultSet resultSet = arrowFlightJdbcArray.getResultSet(3, 5)) { int count = 0; while (resultSet.next()) { - assertEquals((Object) resultSet.getInt(1), dataVector.getObject(count + 3)); + assertEquals((Object) resultSet.getInt(2), dataVector.getObject(count + 3)); count++; } assertEquals(5, count); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java index 1977b61392..7127c7fc32 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionCookieTest.java @@ -39,11 +39,11 @@ public void testCookies() throws SQLException { Statement statement = connection.createStatement()) { // Expect client didn't receive cookies before any operation - assertNull(FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie()); + assertNull(FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie()); // Run another action for check if the cookies was sent by the server. statement.execute(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); - assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getMiddlewareCookieFactory().getCookie()); + assertEquals("k=v", FLIGHT_SERVER_TEST_EXTENSION.getInterceptorFactory().getCookie()); } } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java index ae355829d7..88fb9889b6 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcDriverTest.java @@ -201,6 +201,30 @@ public void testConnectWithInsensitiveCasePropertyKeys2() throws Exception { } } + /** + * Tests whether the {@link ArrowFlightJdbcDriver} can establish a successful connection to the + * Arrow Flight client when provided with null properties. + */ + @Test + public void testConnectWithNullProperties() throws Exception { + final Driver driver = new ArrowFlightJdbcDriver(); + try (Connection connection = + driver.connect( + "jdbc:arrow-flight://" + + dataSource.getConfig().getHost() + + ":" + + dataSource.getConfig().getPort() + + "?" + + "useEncryption=false" + + "&user=" + + dataSource.getConfig().getUser() + + "&password=" + + dataSource.getConfig().getPassword(), + null)) { + assertTrue(connection.isValid(300)); + } + } + /** * Tests whether an exception is thrown upon attempting to connect to a malformed URI. * diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java index 774ad0081e..078837adf3 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java @@ -20,6 +20,8 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.sql.Connection; @@ -83,6 +85,53 @@ public void testSimpleQueryNoParameterBinding() throws SQLException { } } + @Test + public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException { + final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; + try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { + boolean isResultSet = preparedStatement.execute(); + assertTrue(isResultSet); + final ResultSet resultSet = preparedStatement.getResultSet(); + CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet); + assertFalse(preparedStatement.getMoreResults()); + assertEquals(-1, preparedStatement.getUpdateCount()); + } + } + + @Test + public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException { + final String query = "SELECT * FROM TEST_V2"; + final Schema schema = + new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType()))); + PRODUCER.addSelectQuery( + query, + schema, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + root.allocateNew(); + ((IntVector) root.getVector(0)).setSafe(0, 123); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } finally { + listener.completed(); + } + }), + false); + try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { + boolean isResultSet = preparedStatement.execute(); + assertTrue(isResultSet); + final ResultSet resultSet = preparedStatement.getResultSet(); + assertTrue(resultSet.next()); + assertEquals(123, resultSet.getInt(1)); + assertFalse(resultSet.next()); + assertFalse(preparedStatement.getMoreResults()); + assertEquals(-1, preparedStatement.getUpdateCount()); + } + } + @Test public void testQueryWithParameterBinding() throws SQLException { final String query = "Fake query with parameters"; @@ -174,6 +223,34 @@ public void testUpdateQuery() throws SQLException { } } + @Test + public void testUpdateQueryWithExecute() throws SQLException { + String query = "Fake update with execute"; + PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); + try (final PreparedStatement stmt = connection.prepareStatement(query)) { + boolean isResultSet = stmt.execute(); + assertFalse(isResultSet); + int updated = stmt.getUpdateCount(); + assertEquals(42, updated); + assertFalse(stmt.getMoreResults()); + assertEquals(-1, stmt.getUpdateCount()); + } + } + + @Test + public void testUpdateQueryWithExecuteV2() throws SQLException { + String query = "Fake update with execute V2"; + PRODUCER.addUpdateQuery(query, /*updatedRows*/ 99, true); + try (final PreparedStatement stmt = connection.prepareStatement(query)) { + boolean isResultSet = stmt.execute(); + assertFalse(isResultSet); + int updated = stmt.getUpdateCount(); + assertEquals(99, updated); + assertFalse(stmt.getMoreResults()); + assertEquals(-1, stmt.getUpdateCount()); + } + } + @Test public void testUpdateQueryWithParameters() throws SQLException { String query = "Fake update with parameters"; diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java index 632cb0ba56..6acce9c2a6 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java @@ -62,6 +62,9 @@ public class ArrowFlightStatementExecuteTest { private static final String SAMPLE_LARGE_UPDATE_QUERY = "UPDATE this_large_table SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition"; private static final long SAMPLE_LARGE_UPDATE_COUNT = Long.MAX_VALUE; + private static final String SAMPLE_QUERY_CMD_V2 = "SELECT * FROM this_test_v2"; + private static final String SAMPLE_LARGE_UPDATE_QUERY_V2 = + "UPDATE this_large_table_v2 SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition"; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); @RegisterExtension @@ -96,6 +99,31 @@ public static void setUpBeforeClass() { })); PRODUCER.addUpdateQuery(SAMPLE_UPDATE_QUERY, SAMPLE_UPDATE_COUNT); PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY, SAMPLE_LARGE_UPDATE_COUNT); + + // V2 queries with is_update field set + PRODUCER.addSelectQuery( + SAMPLE_QUERY_CMD_V2, + SAMPLE_QUERY_SCHEMA, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = + VectorSchemaRoot.create(SAMPLE_QUERY_SCHEMA, allocator)) { + final UInt1Vector vector = (UInt1Vector) root.getVector(VECTOR_NAME); + IntStream.range(0, SAMPLE_QUERY_ROWS) + .forEach(index -> vector.setSafe(index, index)); + vector.setValueCount(SAMPLE_QUERY_ROWS); + root.setRowCount(SAMPLE_QUERY_ROWS); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + }), + false); + PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY_V2, SAMPLE_LARGE_UPDATE_COUNT, true); } @BeforeEach @@ -168,4 +196,42 @@ public void testUpdateCountShouldStartOnZero() throws SQLException { is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(0L)))); assertThat(statement.getResultSet(), is(nullValue())); } + + @Test + public void testExecuteShouldRunSelectQueryV2() throws SQLException { + assertThat(statement.execute(SAMPLE_QUERY_CMD_V2), is(true)); + final Set numbers = + IntStream.range(0, SAMPLE_QUERY_ROWS) + .boxed() + .map(Integer::byteValue) + .collect(Collectors.toCollection(HashSet::new)); + try (final ResultSet resultSet = statement.getResultSet()) { + final int columnCount = resultSet.getMetaData().getColumnCount(); + assertThat(columnCount, is(1)); + int rowCount = 0; + for (; resultSet.next(); rowCount++) { + assertThat(numbers.remove(resultSet.getByte(1)), is(true)); + } + assertThat(rowCount, is(equalTo(SAMPLE_QUERY_ROWS))); + } + assertThat(numbers, is(Collections.emptySet())); + assertThat( + (long) statement.getUpdateCount(), + is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L)))); + } + + @Test + public void testExecuteShouldRunUpdateQueryForLargeUpdateV2() throws SQLException { + assertThat(statement.execute(SAMPLE_LARGE_UPDATE_QUERY_V2), is(false)); // UPDATE query. + final long updateCountSmall = statement.getUpdateCount(); + final long updateCountLarge = statement.getLargeUpdateCount(); + assertThat(updateCountLarge, is(equalTo(SAMPLE_LARGE_UPDATE_COUNT))); + assertThat( + updateCountSmall, + is( + allOf( + equalTo((long) AvaticaUtils.toSaturatedInt(updateCountLarge)), + not(equalTo(updateCountLarge))))); + assertThat(statement.getResultSet(), is(nullValue())); + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java index 8e872a1167..55722f60fb 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ConnectionTest.java @@ -16,24 +16,42 @@ */ package org.apache.arrow.driver.jdbc; +import static java.lang.String.format; +import static java.util.stream.IntStream.range; +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.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; +import com.google.protobuf.Message; import java.net.URISyntaxException; import java.sql.Connection; import java.sql.Driver; import java.sql.DriverManager; +import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.Statement; +import java.util.Map; import java.util.Properties; +import java.util.function.Consumer; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.flight.FlightMethod; +import org.apache.arrow.flight.FlightProducer.ServerStreamListener; +import org.apache.arrow.flight.NoOpSessionOptionValueVisitor; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -576,4 +594,178 @@ public void testPasswordConnectionPropertyIntegerCorrectCastUrlWithDriverManager assertTrue(connection.isValid(0)); } } + + /** + * Test that the JDBC driver properly integrates driver version into client handler. + * + * @throws Exception on error. + */ + @Test + public void testJdbcDriverVersionIntegration() throws Exception { + final Properties properties = new Properties(); + properties.put( + ArrowFlightConnectionProperty.HOST.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getHost()); + properties.put( + ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); + properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); + properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); + properties.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); + + // Create a driver instance and connect + ArrowFlightJdbcDriver driverVersion = new ArrowFlightJdbcDriver(); + + try (Connection connection = + ArrowFlightConnection.createNewConnection( + driverVersion, + new ArrowFlightJdbcFactory(), + "jdbc:arrow-flight-sql://localhost:" + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties, + allocator)) { + + assertTrue(connection.isValid(0)); + + var actualUserAgent = + FLIGHT_SERVER_TEST_EXTENSION + .getInterceptorFactory() + .getHeader(FlightMethod.HANDSHAKE, "user-agent"); + + var expectedUserAgent = + "JDBC Flight SQL Driver " + driverVersion.getDriverVersion().versionString; + // Driver appends version to grpc user-agent header. Assert the header starts + // with the + // expected + // value and ignored grpc version. + assertTrue( + actualUserAgent.startsWith(expectedUserAgent), + "Expected: " + expectedUserAgent + " but found: " + actualUserAgent); + } + } + + @Test + public void testSetCatalogShouldUpdateSessionOptions() throws Exception { + final Properties properties = new Properties(); + properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); + properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); + properties.put("useEncryption", false); + + try (Connection connection = + DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties)) { + final String catalog = "new_catalog"; + connection.setCatalog(catalog); + + final Map options = PRODUCER.getSessionOptions(); + assertTrue(options.containsKey("catalog")); + String actualCatalog = + options + .get("catalog") + .acceptVisitor( + new NoOpSessionOptionValueVisitor() { + @Override + public String visit(String value) { + return value; + } + }); + assertEquals(catalog, actualCatalog); + } + } + + @Test + public void testStatementsClosedOnConnectionClose() throws Exception { + // create a connection + final Properties properties = new Properties(); + properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); + properties.put( + ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); + properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); + properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); + properties.put("useEncryption", false); + + Connection connection = + DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties); + + // create some statements + int numStatements = 3; + Statement[] statements = new Statement[numStatements]; + for (int i = 0; i < numStatements; i++) { + statements[i] = connection.createStatement(); + assertFalse(statements[i].isClosed()); + } + + // close the connection + connection.close(); + + // assert the statements are closed + for (int i = 0; i < numStatements; i++) { + assertTrue(statements[i].isClosed()); + } + } + + @Test + public void testResultSetsFromDatabaseMetadataClosedOnConnectionClose() throws Exception { + // set up the FlightProducer to respond to metadata queries + // getTableTypes() is being used, but any other method would work + int rowCount = 3; + final Message commandGetTableTypes = CommandGetTableTypes.getDefaultInstance(); + final Consumer commandGetTableTypesResultProducer = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(); + final VectorSchemaRoot root = + VectorSchemaRoot.create(Schemas.GET_TABLE_TYPES_SCHEMA, allocator)) { + final VarCharVector tableType = (VarCharVector) root.getVector("table_type"); + range(0, rowCount) + .forEach(i -> tableType.setSafe(i, new Text(format("table_type #%d", i)))); + root.setRowCount(rowCount); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + }; + PRODUCER.addCatalogQuery(commandGetTableTypes, commandGetTableTypesResultProducer); + + // create a connection + final Properties properties = new Properties(); + properties.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); + properties.put( + ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); + properties.put(ArrowFlightConnectionProperty.USER.camelName(), userTest); + properties.put(ArrowFlightConnectionProperty.PASSWORD.camelName(), passTest); + properties.put("useEncryption", false); + + Connection connection = + DriverManager.getConnection( + "jdbc:arrow-flight-sql://" + + FLIGHT_SERVER_TEST_EXTENSION.getHost() + + ":" + + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + properties); + + // create ResultSets from DatabaseMetadata + int numResultSets = 3; + ResultSet[] resultSets = new ResultSet[numResultSets]; + for (int i = 0; i < numResultSets; i++) { + resultSets[i] = connection.getMetaData().getTableTypes(); + assertFalse(resultSets[i].isClosed()); + } + + // close the connection + connection.close(); + + // assert the ResultSets are closed + for (int i = 0; i < numResultSets; i++) { + assertTrue(resultSets[i].isClosed()); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java index aa586651f5..f71114e1b5 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/FlightServerTestExtension.java @@ -25,6 +25,8 @@ import java.sql.SQLException; import java.util.ArrayDeque; import java.util.Deque; +import java.util.HashMap; +import java.util.Map; import java.util.Properties; import org.apache.arrow.driver.jdbc.authentication.Authentication; import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication; @@ -33,6 +35,7 @@ import org.apache.arrow.flight.CallHeaders; import org.apache.arrow.flight.CallInfo; import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightMethod; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightServerMiddleware; import org.apache.arrow.flight.Location; @@ -67,7 +70,8 @@ public class FlightServerTestExtension private final CertKeyPair certKeyPair; private final File mTlsCACert; - private final MiddlewareCookie.Factory middlewareCookieFactory = new MiddlewareCookie.Factory(); + private final InterceptorMiddleware.Factory interceptorFactory = + new InterceptorMiddleware.Factory(); private FlightServerTestExtension( final Properties properties, @@ -126,12 +130,18 @@ public Connection getConnection(boolean useEncryption) throws SQLException { return this.createDataSource().getConnection(); } + public Connection getConnection(String timezone) throws SQLException { + setUseEncryption(false); + properties.put("timezone", timezone); + return this.createDataSource().getConnection(); + } + private void setUseEncryption(boolean useEncryption) { properties.put("useEncryption", useEncryption); } - public MiddlewareCookie.Factory getMiddlewareCookieFactory() { - return middlewareCookieFactory; + public InterceptorMiddleware.Factory getInterceptorFactory() { + return interceptorFactory; } @FunctionalInterface @@ -143,7 +153,7 @@ private FlightServer initiateServer(Location location) throws IOException { FlightServer.Builder builder = FlightServer.builder(allocator, location, producer) .headerAuthenticator(authentication.authenticate()) - .middleware(FlightServerMiddleware.Key.of("KEY"), middlewareCookieFactory); + .middleware(FlightServerMiddleware.Key.of("KEY"), interceptorFactory); if (certKeyPair != null) { builder.useTls(certKeyPair.cert, certKeyPair.key); } @@ -301,11 +311,11 @@ public FlightServerTestExtension build() { * A middleware to handle with the cookies in the server. It is used to test if cookies are being * sent properly. */ - static class MiddlewareCookie implements FlightServerMiddleware { + static class InterceptorMiddleware implements FlightServerMiddleware { private final Factory factory; - public MiddlewareCookie(Factory factory) { + public InterceptorMiddleware(Factory factory) { this.factory = factory; } @@ -323,22 +333,33 @@ public void onCallCompleted(CallStatus callStatus) {} public void onCallErrored(Throwable throwable) {} /** A factory for the MiddlewareCookie. */ - static class Factory implements FlightServerMiddleware.Factory { + static class Factory implements FlightServerMiddleware.Factory { + private final Map receivedCallHeaders = new HashMap<>(); private boolean receivedCookieHeader = false; private String cookie; @Override - public MiddlewareCookie onCallStarted( + public InterceptorMiddleware onCallStarted( CallInfo callInfo, CallHeaders callHeaders, RequestContext requestContext) { cookie = callHeaders.get("Cookie"); receivedCookieHeader = null != cookie; - return new MiddlewareCookie(this); + + receivedCallHeaders.put(callInfo.method(), callHeaders); + return new InterceptorMiddleware(this); } public String getCookie() { return cookie; } + + public String getHeader(FlightMethod method, String key) { + CallHeaders headers = receivedCallHeaders.get(method); + if (headers == null) { + return null; + } + return headers.get(key); + } } } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java new file mode 100644 index 0000000000..5e782db031 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/OAuthIntegrationTest.java @@ -0,0 +1,474 @@ +/* + * 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.arrow.driver.jdbc; + +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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.net.URI; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.util.Properties; +import java.util.concurrent.TimeUnit; +import mockwebserver3.MockResponse; +import mockwebserver3.MockWebServer; +import mockwebserver3.RecordedRequest; +import mockwebserver3.junit5.StartStop; +import org.apache.arrow.driver.jdbc.authentication.TokenAuthentication; +import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty; +import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Integration tests for OAuth authentication flows in the JDBC driver. + * + *

These tests verify that OAuth tokens obtained from an OAuth server are correctly used in + * Flight SQL requests. + */ +public class OAuthIntegrationTest { + + private static final String VALID_ACCESS_TOKEN = "valid-oauth-access-token-12345"; + private static final String CLIENT_ID = "test-client-id"; + private static final String CLIENT_SECRET = "test-client-secret"; + private static final String SUBJECT_TOKEN = "original-subject-token"; + private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"; + private static final String TEST_SCOPE = "dremio.all"; + + private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer(); + + @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; + + static { + FLIGHT_SERVER_TEST_EXTENSION = + new FlightServerTestExtension.Builder() + .authentication(new TokenAuthentication.Builder().token(VALID_ACCESS_TOKEN).build()) + .producer(FLIGHT_SQL_PRODUCER) + .build(); + } + + @StartStop private final MockWebServer oauthServer = new MockWebServer(); + private URI tokenEndpoint; + + @BeforeAll + public static void setUpClass() { + // Register a simple catalog query handler + FLIGHT_SQL_PRODUCER.addCatalogQuery( + CommandGetCatalogs.getDefaultInstance(), + listener -> { + try (BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = + VectorSchemaRoot.create(Schemas.GET_CATALOGS_SCHEMA, allocator)) { + root.setRowCount(0); + listener.start(root); + listener.putNext(); + } catch (Throwable t) { + listener.error(t); + } finally { + listener.completed(); + } + }); + + // Register a simple schema query handler for getSchemas() + FLIGHT_SQL_PRODUCER.addCatalogQuery( + CommandGetDbSchemas.getDefaultInstance(), + listener -> { + try (BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = + VectorSchemaRoot.create(Schemas.GET_SCHEMAS_SCHEMA, allocator)) { + root.setRowCount(0); + listener.start(root); + listener.putNext(); + } catch (Throwable t) { + listener.error(t); + } finally { + listener.completed(); + } + }); + } + + @AfterAll + public static void tearDownClass() { + AutoCloseables.closeNoChecked(FLIGHT_SQL_PRODUCER); + } + + @BeforeEach + public void setUp() { + tokenEndpoint = oauthServer.url("/oauth/token").uri(); + } + + @AfterEach + public void tearDown() { + oauthServer.close(); + } + + // Helper methods for mock OAuth responses + + private void enqueueSuccessfulTokenResponse() { + enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600); + } + + private void enqueueSuccessfulTokenResponse(String token, int expiresIn) { + String body = + String.format( + "{\"access_token\":\"%s\",\"token_type\":\"Bearer\",\"expires_in\":%d}", + token, expiresIn); + oauthServer.enqueue( + new MockResponse.Builder() + .code(200) + .setHeader("Content-Type", "application/json") + .body(body) + .build()); + } + + private void enqueueErrorResponse(String error, String description) { + String body = + String.format("{\"error\":\"%s\",\"error_description\":\"%s\"}", error, description); + oauthServer.enqueue( + new MockResponse.Builder() + .code(400) + .setHeader("Content-Type", "application/json") + .body(body) + .build()); + } + + private Properties createBaseProperties() { + Properties props = new Properties(); + props.put(ArrowFlightConnectionProperty.HOST.camelName(), "localhost"); + props.put( + ArrowFlightConnectionProperty.PORT.camelName(), FLIGHT_SERVER_TEST_EXTENSION.getPort()); + props.put(ArrowFlightConnectionProperty.USE_ENCRYPTION.camelName(), false); + return props; + } + + private String getJdbcUrl() { + return String.format( + "jdbc:arrow-flight-sql://localhost:%d", FLIGHT_SERVER_TEST_EXTENSION.getPort()); + } + + // ==================== Client Credentials Flow Tests ==================== + + @Test + public void testClientCredentialsFlowSuccess() throws Exception { + enqueueSuccessfulTokenResponse(); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + assertFalse(conn.isClosed()); + // Trigger a Flight call to force OAuth token retrieval + conn.getMetaData().getCatalogs().close(); + } + + // Verify OAuth request was made + RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(oauthRequest, "OAuth request should have been made"); + assertEquals("POST", oauthRequest.getMethod()); + String body = oauthRequest.getBody().utf8(); + assertTrue(body.contains("grant_type=client_credentials")); + assertTrue(body.contains("scope=" + TEST_SCOPE)); + } + + @Test + public void testClientCredentialsFlowWithUrlParameters() throws Exception { + enqueueSuccessfulTokenResponse(); + + String url = + String.format( + "jdbc:arrow-flight-sql://localhost:%d?useEncryption=false" + + "&oauth.flow=client_credentials" + + "&oauth.tokenUri=%s" + + "&oauth.clientId=%s" + + "&oauth.clientSecret=%s", + FLIGHT_SERVER_TEST_EXTENSION.getPort(), + tokenEndpoint.toString(), + CLIENT_ID, + CLIENT_SECRET); + + try (Connection conn = DriverManager.getConnection(url)) { + conn.getMetaData().getCatalogs().close(); + } + + assertEquals(1, oauthServer.getRequestCount()); + } + + @Test + public void testClientCredentialsFlowInvalidCredentials() throws Exception { + enqueueErrorResponse("invalid_client", "Client authentication failed"); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), "wrong-client"); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), "wrong-secret"); + + Exception ex = + assertThrows( + Exception.class, + () -> { + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + conn.getMetaData().getCatalogs().close(); + } + }); + // Verify the error message contains the OAuth error somewhere in the exception chain + assertTrue( + containsInExceptionChain(ex, "invalid_client"), + "Exception chain should contain 'invalid_client'"); + } + + private boolean containsInExceptionChain(Throwable t, String message) { + while (t != null) { + if (t.getMessage() != null && t.getMessage().contains(message)) { + return true; + } + t = t.getCause(); + } + return false; + } + + // ==================== Token Exchange Flow Tests ==================== + + @Test + public void testTokenExchangeFlowMinimalParameters() throws Exception { + enqueueSuccessfulTokenResponse(); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(), + SUBJECT_TOKEN_TYPE); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + conn.getMetaData().getCatalogs().close(); + } + + RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(oauthRequest, "OAuth request should have been made"); + String body = oauthRequest.getBody().utf8(); + assertTrue( + body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange"), + "Should contain token exchange grant type"); + assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN)); + } + + @Test + public void testTokenExchangeFlowWithAllParameters() throws Exception { + enqueueSuccessfulTokenResponse(); + + String actorToken = "actor-token-value"; + String actorTokenType = "urn:ietf:params:oauth:token-type:access_token"; + String audience = "https://api.example.com"; + String resource = "https://api.example.com/resource"; + String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token"; + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + props.put(ArrowFlightConnectionProperty.OAUTH_SCOPE.camelName(), TEST_SCOPE); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(), + SUBJECT_TOKEN_TYPE); + props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN.camelName(), actorToken); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_ACTOR_TOKEN_TYPE.camelName(), actorTokenType); + props.put(ArrowFlightConnectionProperty.OAUTH_EXCHANGE_AUDIENCE.camelName(), audience); + props.put(ArrowFlightConnectionProperty.OAUTH_RESOURCE.camelName(), resource); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.camelName(), + requestedTokenType); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + conn.getMetaData().getCatalogs().close(); + } + + RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(oauthRequest, "OAuth request should have been made"); + String body = oauthRequest.getBody().utf8(); + assertTrue(body.contains("subject_token=" + SUBJECT_TOKEN)); + assertTrue(body.contains("actor_token=" + actorToken)); + } + + @Test + public void testTokenExchangeFlowWithClientAuthentication() throws Exception { + enqueueSuccessfulTokenResponse(); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(), + SUBJECT_TOKEN_TYPE); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + conn.getMetaData().getCatalogs().close(); + } + + RecordedRequest oauthRequest = oauthServer.takeRequest(5, TimeUnit.SECONDS); + assertNotNull(oauthRequest, "OAuth request should have been made"); + String authHeader = oauthRequest.getHeaders().get("Authorization"); + assertNotNull(authHeader, "Should have Basic auth header for client authentication"); + assertTrue(authHeader.startsWith("Basic ")); + } + + // ==================== Token Caching Tests ==================== + + @Test + public void testTokenCachingAcrossMultipleOperations() throws Exception { + enqueueSuccessfulTokenResponse(); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + // Execute multiple operations + conn.isValid(5); + conn.getMetaData().getCatalogs().close(); + conn.getMetaData().getSchemas().close(); + } + + // Should only have made one OAuth request due to caching + assertEquals(1, oauthServer.getRequestCount()); + } + + @Test + public void testTokenRefreshAfterExpiration() throws Exception { + enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 1); + enqueueSuccessfulTokenResponse(VALID_ACCESS_TOKEN, 3600); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN.camelName(), SUBJECT_TOKEN); + props.put( + ArrowFlightConnectionProperty.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.camelName(), + SUBJECT_TOKEN_TYPE); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + // First operation triggers initial token fetch + conn.getMetaData().getCatalogs().close(); + + // Token with 1s expiry is immediately considered expired (due to 30s buffer) + // so the next operation should trigger a refresh + conn.getMetaData().getCatalogs().close(); + } + + // Should have made exactly 2 OAuth requests: initial + refresh + assertEquals(2, oauthServer.getRequestCount()); + } + + // ==================== Error Handling Tests ==================== + + @Test + public void testMissingRequiredParametersClientCredentials() { + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + // Missing client_id and client_secret + + assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props)); + } + + @Test + public void testMissingRequiredParametersTokenExchange() { + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "token_exchange"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + // Missing subject_token and subject_token_type + + assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props)); + } + + @Test + public void testInvalidOAuthFlow() { + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "invalid_flow"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + + assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props)); + } + + @Test + public void testMalformedTokenEndpoint() { + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), "not-a-valid-uri://"); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + + assertThrows(SQLException.class, () -> DriverManager.getConnection(getJdbcUrl(), props)); + } + + // ==================== Authorization Header Verification ==================== + + @Test + public void testOAuthTokenSentAsBearer() throws Exception { + enqueueSuccessfulTokenResponse(); + + Properties props = createBaseProperties(); + props.put(ArrowFlightConnectionProperty.OAUTH_FLOW.camelName(), "client_credentials"); + props.put(ArrowFlightConnectionProperty.OAUTH_TOKEN_URI.camelName(), tokenEndpoint.toString()); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_ID.camelName(), CLIENT_ID); + props.put(ArrowFlightConnectionProperty.OAUTH_CLIENT_SECRET.camelName(), CLIENT_SECRET); + + try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) { + conn.getMetaData().getCatalogs().close(); + } + + // Verify the Flight server received the bearer token + String authHeader = + FLIGHT_SERVER_TEST_EXTENSION + .getInterceptorFactory() + .getHeader(org.apache.arrow.flight.FlightMethod.GET_FLIGHT_INFO, "authorization"); + assertNotNull(authHeader, "Authorization header should be present in Flight requests"); + assertEquals("Bearer " + VALID_ACCESS_TOKEN, authHeader); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index cd47408f52..3a5a39be3d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -22,8 +22,10 @@ import static org.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.anyOf; import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.*; @@ -31,7 +33,9 @@ import java.nio.charset.StandardCharsets; import java.sql.Connection; import java.sql.DriverManager; +import java.sql.PreparedStatement; import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.sql.Statement; @@ -42,6 +46,7 @@ import java.util.List; import java.util.Random; import java.util.Set; +import java.util.UUID; import java.util.concurrent.CountDownLatch; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.FallbackFlightSqlProducer; @@ -61,6 +66,7 @@ import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -699,10 +705,9 @@ public void testFallbackUnresolvableFlightServer() throws Exception { } attempt1 = System.nanoTime(); elapsedMs = (attempt1 - start) / 1_000_000.; - // TODO(GH-661): this assertion should be flipped to assertTrue. - assertFalse( + assertTrue( elapsedMs < 5000., - String.format("Expected second attempt to be the same, but %f ms elapsed", elapsedMs)); + String.format("Expected second attempt to be faster, but %f ms elapsed", elapsedMs)); } } } @@ -796,4 +801,174 @@ public void testResultSetAppMetadata() throws Exception { "foo".getBytes(StandardCharsets.UTF_8)); } } + + @Test + public void testSelectQueryWithUuidColumn() throws SQLException { + // Expectations + final int expectedRowCount = 4; + final UUID[] expectedUuids = + new UUID[] { + CoreMockedSqlProducers.UUID_1, + CoreMockedSqlProducers.UUID_2, + CoreMockedSqlProducers.UUID_3, + null + }; + + final Integer[] expectedIds = new Integer[] {1, 2, 3, 4}; + + final List actualUuids = new ArrayList<>(expectedRowCount); + final List actualIds = new ArrayList<>(expectedRowCount); + + // Query + int actualRowCount = 0; + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + for (; resultSet.next(); actualRowCount++) { + actualIds.add((Integer) resultSet.getObject("id")); + actualUuids.add((UUID) resultSet.getObject("uuid_col")); + } + } + + // Assertions + int finalActualRowCount = actualRowCount; + assertAll( + "UUID ResultSet values are as expected", + () -> assertThat(finalActualRowCount, is(equalTo(expectedRowCount))), + () -> assertThat(actualIds.toArray(new Integer[0]), is(expectedIds)), + () -> assertThat(actualUuids.toArray(new UUID[0]), is(expectedUuids))); + } + + @Test + public void testGetObjectReturnsUuid() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + Object result = resultSet.getObject("uuid_col"); + assertThat(result, instanceOf(UUID.class)); + assertThat(result, is(CoreMockedSqlProducers.UUID_1)); + } + } + + @Test + public void testGetObjectByIndexReturnsUuid() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + Object result = resultSet.getObject(2); + assertThat(result, instanceOf(UUID.class)); + assertThat(result, is(CoreMockedSqlProducers.UUID_1)); + } + } + + @Test + public void testGetStringReturnsHyphenatedFormat() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + String result = resultSet.getString("uuid_col"); + assertThat(result, is(CoreMockedSqlProducers.UUID_1.toString())); + } + } + + @Test + public void testGetBytesReturns16ByteArray() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + byte[] result = resultSet.getBytes("uuid_col"); + assertThat(result.length, is(16)); + assertThat(result, is(UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1))); + } + } + + @Test + public void testNullUuidHandling() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + // Skip to row 4 which has NULL UUID + resultSet.next(); // row 1 + resultSet.next(); // row 2 + resultSet.next(); // row 3 + resultSet.next(); // row 4 (NULL UUID) + + Object objResult = resultSet.getObject("uuid_col"); + assertThat(objResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + + String strResult = resultSet.getString("uuid_col"); + assertThat(strResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + + byte[] bytesResult = resultSet.getBytes("uuid_col"); + assertThat(bytesResult, nullValue()); + assertThat(resultSet.wasNull(), is(true)); + } + } + + @Test + public void testMultipleUuidRows() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_2)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_3)); + + resultSet.next(); + assertThat(resultSet.getObject("uuid_col"), nullValue()); + } + } + + @Test + public void testUuidExtensionTypeInSchema() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.UUID_SQL_CMD)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + + assertThat(metaData.getColumnCount(), is(2)); + assertThat(metaData.getColumnName(1), is("id")); + assertThat(metaData.getColumnName(2), is("uuid_col")); + + assertThat(metaData.getColumnType(2), is(java.sql.Types.OTHER)); + } + } + + @Test + public void testPreparedStatementWithUuidParameter() throws SQLException { + try (PreparedStatement pstmt = + connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) { + pstmt.setObject(1, CoreMockedSqlProducers.UUID_1); + try (ResultSet rs = pstmt.executeQuery()) { + rs.next(); + assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + } + } + } + + @Test + public void testPreparedStatementWithUuidStringParameter() throws SQLException { + try (PreparedStatement pstmt = + connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) { + pstmt.setString(1, CoreMockedSqlProducers.UUID_1.toString()); + try (ResultSet rs = pstmt.executeQuery()) { + rs.next(); + assertThat(rs.getObject("uuid_col"), is(CoreMockedSqlProducers.UUID_1)); + } + } + } + + @Test + public void testPreparedStatementUpdateWithUuid() throws SQLException { + try (PreparedStatement pstmt = + connection.prepareStatement(CoreMockedSqlProducers.UUID_PREPARED_UPDATE_SQL_CMD)) { + pstmt.setObject(1, CoreMockedSqlProducers.UUID_3); + pstmt.setInt(2, 1); + int updated = pstmt.executeUpdate(); + assertThat(updated, is(1)); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java new file mode 100644 index 0000000000..0921ae2d38 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/TimestampResultSetTest.java @@ -0,0 +1,164 @@ +/* + * 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.arrow.driver.jdbc; + +import com.google.common.collect.ImmutableList; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.sql.Types; +import java.time.Instant; +import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.Calendar; +import java.util.Collections; +import java.util.TimeZone; +import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.TimeStampVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Timestamps have a lot of nuances in JDBC. This class is here to test that timestamp behavior is + * correct for different types of Timestamp vectors as well as different methods of retrieving the + * timestamps in JDBC. + */ +public class TimestampResultSetTest { + private static final MockFlightSqlProducer FLIGHT_SQL_PRODUCER = new MockFlightSqlProducer(); + + @RegisterExtension public static FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; + + static { + FLIGHT_SERVER_TEST_EXTENSION = + FlightServerTestExtension.createStandardTestExtension(FLIGHT_SQL_PRODUCER); + } + + private static final String QUERY_STRING = "SELECT * FROM TIMESTAMPS"; + private static final Schema QUERY_SCHEMA = + new Schema( + ImmutableList.of( + Field.nullable("no_tz", new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + Field.nullable("utc", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + Field.nullable("utc+1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT+1")), + Field.nullable("utc-1", new ArrowType.Timestamp(TimeUnit.MILLISECOND, "GMT-1")))); + + @BeforeAll + public static void setup() throws SQLException { + Instant firstDay2025 = OffsetDateTime.of(2025, 1, 1, 0, 0, 0, 0, ZoneOffset.UTC).toInstant(); + + FLIGHT_SQL_PRODUCER.addSelectQuery( + QUERY_STRING, + QUERY_SCHEMA, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) { + listener.start(root); + root.getFieldVectors() + .forEach(v -> ((TimeStampVector) v).setSafe(0, firstDay2025.toEpochMilli())); + root.setRowCount(1); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); + } + + /** + * This test doesn't yet test anything other than ensuring all ResultSet methods to retrieve a + * timestamp succeed. + * + *

This is a good starting point to add more tests to ensure the values are correct when we + * change the "local calendar" either through changing the JVM default or through the connection + * property. + */ + @Test + public void test() { + TimeZone.setDefault(TimeZone.getTimeZone("UTC")); + try (Connection connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection("UTC")) { + try (PreparedStatement s = connection.prepareStatement(QUERY_STRING)) { + try (ResultSet rs = s.executeQuery()) { + int numCols = rs.getMetaData().getColumnCount(); + try { + rs.next(); + for (int i = 1; i <= numCols; i++) { + int type = rs.getMetaData().getColumnType(i); + String name = rs.getMetaData().getColumnName(i); + System.out.println(name); + System.out.print("- getDate:\t\t\t\t\t\t\t"); + System.out.print(rs.getDate(i)); + System.out.println(); + System.out.print("- getTimestamp:\t\t\t\t\t\t"); + System.out.print(rs.getTimestamp(i)); + System.out.println(); + System.out.print("- getString:\t\t\t\t\t\t"); + System.out.print(rs.getString(i)); + System.out.println(); + System.out.print("- getObject:\t\t\t\t\t\t"); + System.out.print(rs.getObject(i)); + System.out.println(); + System.out.print("- getObject(Timestamp.class):\t\t"); + System.out.print(rs.getObject(i, Timestamp.class)); + System.out.println(); + System.out.print("- getTimestamp(default Calendar):\t"); + System.out.print(rs.getTimestamp(i, Calendar.getInstance())); + System.out.println(); + System.out.print("- getTimestamp(UTC Calendar):\t\t"); + System.out.print( + rs.getTimestamp(i, Calendar.getInstance(TimeZone.getTimeZone("UTC")))); + System.out.println(); + System.out.print("- getObject(LocalDateTime.class):\t"); + System.out.print(rs.getObject(i, LocalDateTime.class)); + System.out.println(); + if (type == Types.TIMESTAMP_WITH_TIMEZONE) { + System.out.print("- getObject(Instant.class):\t\t\t"); + System.out.print(rs.getObject(i, Instant.class)); + System.out.println(); + System.out.print("- getObject(OffsetDateTime.class):\t"); + System.out.print(rs.getObject(i, OffsetDateTime.class)); + System.out.println(); + System.out.print("- getObject(ZonedDateTime.class):\t"); + System.out.print(rs.getObject(i, ZonedDateTime.class)); + System.out.println(); + } + System.out.println(); + } + System.out.println(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java index b56bf3c63d..1fbd2f86a9 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/ArrowFlightJdbcAccessorFactoryTest.java @@ -16,10 +16,12 @@ */ package org.apache.arrow.driver.jdbc.accessor; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.function.IntSupplier; import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcBinaryVectorAccessor; +import org.apache.arrow.driver.jdbc.accessor.impl.binary.ArrowFlightJdbcUuidVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDateVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcDurationVectorAccessor; import org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcIntervalVectorAccessor; @@ -46,6 +48,8 @@ import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarBinaryVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.complex.StructVector; @@ -239,6 +243,18 @@ public void createAccessorForFixedSizeBinaryVector() { } } + @Test + public void createAccessorForViewVarBinaryVector() { + try (ValueVector valueVector = + new ViewVarBinaryVector("", rootAllocatorTestExtension.getRootAllocator())) { + ArrowFlightJdbcAccessor accessor = + ArrowFlightJdbcAccessorFactory.createAccessor( + valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); + + assertTrue(accessor instanceof ArrowFlightJdbcBinaryVectorAccessor); + } + } + @Test public void createAccessorForTimeStampVector() { try (ValueVector valueVector = rootAllocatorTestExtension.createTimeStampMilliVector()) { @@ -340,6 +356,18 @@ public void createAccessorForLargeVarCharVector() { } } + @Test + public void createAccessorForViewVarCharVector() { + try (ValueVector valueVector = + new ViewVarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { + ArrowFlightJdbcAccessor accessor = + ArrowFlightJdbcAccessorFactory.createAccessor( + valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); + + assertTrue(accessor instanceof ArrowFlightJdbcVarCharVectorAccessor); + } + } + @Test public void createAccessorForDurationVector() { try (ValueVector valueVector = @@ -471,4 +499,15 @@ public void createAccessorForMapVector() { assertTrue(accessor instanceof ArrowFlightJdbcMapVectorAccessor); } } + + @Test + public void createAccessorForUuidVector() { + try (ValueVector valueVector = rootAllocatorTestExtension.createUuidVector()) { + ArrowFlightJdbcAccessor accessor = + ArrowFlightJdbcAccessorFactory.createAccessor( + valueVector, GET_CURRENT_ROW, (boolean wasNull) -> {}); + + assertInstanceOf(ArrowFlightJdbcUuidVectorAccessor.class, accessor); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java new file mode 100644 index 0000000000..b7f341240c --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/binary/ArrowFlightJdbcUuidVectorAccessorTest.java @@ -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.arrow.driver.jdbc.accessor.impl.binary; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.UUID; +import org.apache.arrow.driver.jdbc.accessor.ArrowFlightJdbcAccessorFactory; +import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.util.UuidUtility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Tests for {@link ArrowFlightJdbcUuidVectorAccessor}. + * + *

Verifies that the accessor correctly handles UUID values from Arrow's UUID extension type, + * following PostgreSQL JDBC driver conventions. + */ +public class ArrowFlightJdbcUuidVectorAccessorTest { + + @RegisterExtension + public static RootAllocatorTestExtension rootAllocatorTestExtension = + new RootAllocatorTestExtension(); + + private static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + private static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + private static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + private UuidVector vector; + private ArrowFlightJdbcUuidVectorAccessor accessor; + private boolean wasNullCalled; + private boolean wasNullValue; + + @BeforeEach + public void setUp() { + vector = rootAllocatorTestExtension.createUuidVector(); + wasNullCalled = false; + wasNullValue = false; + ArrowFlightJdbcAccessorFactory.WasNullConsumer wasNullConsumer = + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }; + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, wasNullConsumer); + } + + @AfterEach + public void tearDown() { + vector.close(); + } + + @Test + public void testGetObjectReturnsUuid() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + Object result = accessor.getObject(); + assertThat(result, is(UUID_1)); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetObjectReturnsCorrectUuidForEachRow() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_1)); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 1, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_2)); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 2, (wasNull) -> {}); + assertThat(accessor.getObject(), is(UUID_3)); + } + + @Test + public void testGetObjectReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + Object result = accessor.getObject(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testGetObjectClassReturnsUuidClass() { + assertThat(accessor.getObjectClass(), equalTo(UUID.class)); + } + + @Test + public void testGetStringReturnsHyphenatedFormat() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + String result = accessor.getString(); + assertThat(result, is("550e8400-e29b-41d4-a716-446655440000")); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetStringReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + String result = accessor.getString(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testGetBytesReturns16ByteArray() { + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + byte[] result = accessor.getBytes(); + assertThat(result.length, is(16)); + assertThat(result, is(UuidUtility.getBytesFromUUID(UUID_1))); + assertThat(accessor.wasNull(), is(false)); + } + + @Test + public void testGetBytesReturnsNullForNullValue() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = new ArrowFlightJdbcUuidVectorAccessor(vector, () -> 0, (wasNull) -> {}); + byte[] result = accessor.getBytes(); + assertThat(result, nullValue()); + assertThat(accessor.wasNull(), is(true)); + } + + @Test + public void testWasNullConsumerIsCalled() { + accessor = + new ArrowFlightJdbcUuidVectorAccessor( + vector, + () -> 0, + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }); + accessor.getObject(); + assertThat(wasNullCalled, is(true)); + assertThat(wasNullValue, is(false)); + } + + @Test + public void testWasNullConsumerIsCalledWithTrueForNull() { + vector.reset(); + vector.allocateNew(1); + vector.setNull(0); + vector.setValueCount(1); + + accessor = + new ArrowFlightJdbcUuidVectorAccessor( + vector, + () -> 0, + (wasNull) -> { + wasNullCalled = true; + wasNullValue = wasNull; + }); + accessor.getObject(); + assertThat(wasNullCalled, is(true)); + assertThat(wasNullValue, is(true)); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java index ad689837e2..c5eb6e34ef 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/AbstractArrowFlightJdbcListAccessorTest.java @@ -191,7 +191,12 @@ public void testShouldGetArrayGetResultSetReturnValidResultSet( try (ResultSet rs = array.getResultSet()) { int count = 0; while (rs.next()) { - final int value = rs.getInt(1); + // Column 1: 1-based index (per JDBC spec) + final int index = rs.getInt(1); + assertThat(index, equalTo(count + 1)); + + // Column 2: actual value (per JDBC spec) + final int value = rs.getInt(2); assertThat(value, equalTo(currentRow * count)); count++; } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java index 696e5afb71..f2d1725fd8 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/complex/ArrowFlightJdbcMapVectorAccessorTest.java @@ -153,15 +153,15 @@ public void testShouldGetArrayReturnValidArray() throws SQLException { try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); - Map entry = resultSet.getObject(1, Map.class); + Map entry = resultSet.getObject(2, Map.class); assertEquals(1, entry.get("key")); assertEquals(11, entry.get("value")); assertTrue(resultSet.next()); - entry = resultSet.getObject(1, Map.class); + entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertEquals(22, entry.get("value")); assertTrue(resultSet.next()); - entry = resultSet.getObject(1, Map.class); + entry = resultSet.getObject(2, Map.class); assertEquals(3, entry.get("key")); assertEquals(33, entry.get("value")); assertFalse(resultSet.next()); @@ -173,7 +173,7 @@ public void testShouldGetArrayReturnValidArray() throws SQLException { assertFalse(accessor.wasNull()); try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); - Map entry = resultSet.getObject(1, Map.class); + Map entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertNull(entry.get("value")); assertFalse(resultSet.next()); @@ -185,19 +185,19 @@ public void testShouldGetArrayReturnValidArray() throws SQLException { assertFalse(accessor.wasNull()); try (ResultSet resultSet = array.getResultSet()) { assertTrue(resultSet.next()); - Map entry = resultSet.getObject(1, Map.class); + Map entry = resultSet.getObject(2, Map.class); assertEquals(0, entry.get("key")); assertEquals(2000, entry.get("value")); assertTrue(resultSet.next()); - entry = resultSet.getObject(1, Map.class); + entry = resultSet.getObject(2, Map.class); assertEquals(1, entry.get("key")); assertEquals(2001, entry.get("value")); assertTrue(resultSet.next()); - entry = resultSet.getObject(1, Map.class); + entry = resultSet.getObject(2, Map.class); assertEquals(2, entry.get("key")); assertEquals(2002, entry.get("value")); assertTrue(resultSet.next()); - entry = resultSet.getObject(1, Map.class); + entry = resultSet.getObject(2, Map.class); assertEquals(3, entry.get("key")); assertEquals(2003, entry.get("value")); assertFalse(resultSet.next()); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java index a2f6fd586f..82876f4aa1 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/text/ArrowFlightJdbcVarCharVectorAccessorTest.java @@ -24,6 +24,8 @@ import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.when; @@ -46,6 +48,8 @@ import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.TimeStampVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -695,4 +699,26 @@ public void testShouldGetObjectClassReturnString() { final Class clazz = accessor.getObjectClass(); assertThat(clazz, equalTo(String.class)); } + + @Test + public void testViewVarcharVector() throws Exception { + try (VarCharVector varCharVector = + new VarCharVector("", rootAllocatorTestExtension.getRootAllocator()); + ViewVarCharVector viewVarCharVector = + new ViewVarCharVector("", rootAllocatorTestExtension.getRootAllocator())) { + varCharVector.allocateNew(1); + viewVarCharVector.allocateNew(1); + + ArrowFlightJdbcVarCharVectorAccessor varCharVectorAccessor = + new ArrowFlightJdbcVarCharVectorAccessor(varCharVector, () -> 0, (boolean wasNull) -> {}); + ArrowFlightJdbcVarCharVectorAccessor viewVarcharVectorAccessor = + new ArrowFlightJdbcVarCharVectorAccessor( + viewVarCharVector, () -> 0, (boolean wasNull) -> {}); + assertNull(viewVarcharVectorAccessor.getString()); + + varCharVector.set(0, new Text("looooong_string")); + viewVarCharVector.set(0, new Text("looooong_string")); + assertEquals(varCharVectorAccessor.getString(), viewVarcharVectorAccessor.getString()); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java index 7b416638e1..a60a71f23d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java @@ -147,7 +147,9 @@ public void testDefaults() { assertNull(builder.clientCertificatePath); assertNull(builder.clientKeyPath); assertEquals(Optional.empty(), builder.catalog); + assertNull(builder.flightClientCache); assertNull(builder.connectTimeout); + assertNull(builder.driverVersion); } @Test diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java new file mode 100644 index 0000000000..d5973ab5d8 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerTest.java @@ -0,0 +1,88 @@ +/* + * 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.arrow.driver.jdbc.client; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; + +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Optional; +import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.CloseSessionRequest; +import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.flight.sql.FlightSqlClient; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +public class ArrowFlightSqlClientHandlerTest { + + @ParameterizedTest + @MethodSource + public void testCloseHandlesFlightRuntimeException( + boolean throwFromCloseSession, CallStatus callStatus, boolean shouldSuppress) + throws Exception { + FlightSqlClient sqlClient = mock(FlightSqlClient.class); + String cacheKey = "cacheKey"; + Optional catalog = + throwFromCloseSession ? Optional.of("test_catalog") : Optional.empty(); + final Collection credentialOptions = new ArrayList<>(); + ArrowFlightSqlClientHandler.Builder builder = new ArrowFlightSqlClientHandler.Builder(); + + if (throwFromCloseSession) { + doThrow(callStatus.toRuntimeException()) + .when(sqlClient) + .closeSession(any(CloseSessionRequest.class), any(CallOption[].class)); + } else { + doThrow(callStatus.toRuntimeException()).when(sqlClient).close(); + } + + ArrowFlightSqlClientHandler sqlClientHandler = + new ArrowFlightSqlClientHandler( + cacheKey, sqlClient, builder, credentialOptions, catalog, null); + + if (shouldSuppress) { + assertDoesNotThrow(sqlClientHandler::close); + } else { + assertThrows(SQLException.class, sqlClientHandler::close); + } + } + + private static Object[] testCloseHandlesFlightRuntimeException() { + CallStatus benignInternalError = + new CallStatus(FlightStatusCode.INTERNAL, null, "Connection closed after GOAWAY", null); + CallStatus notBenignInternalError = + new CallStatus(FlightStatusCode.INTERNAL, null, "Not a benign internal error", null); + CallStatus unavailableError = new CallStatus(FlightStatusCode.UNAVAILABLE, null, null, null); + CallStatus unknownError = new CallStatus(FlightStatusCode.UNKNOWN, null, null, null); + return new Object[] { + new Object[] {true, benignInternalError, true}, + new Object[] {false, benignInternalError, true}, + new Object[] {true, notBenignInternalError, false}, + new Object[] {false, notBenignInternalError, false}, + new Object[] {true, unavailableError, true}, + new Object[] {false, unavailableError, true}, + new Object[] {true, unknownError, false}, + new Object[] {false, unknownError, false}, + }; + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java new file mode 100644 index 0000000000..c258a7c652 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthConfigurationTest.java @@ -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. + */ +package org.apache.arrow.driver.jdbc.client.oauth; + +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.assertThrows; + +import com.nimbusds.oauth2.sdk.Scope; +import java.net.URI; +import java.sql.SQLException; +import java.util.Collections; +import java.util.stream.Stream; +import org.junit.jupiter.api.Named; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** Tests for {@link OAuthConfiguration}. */ +public class OAuthConfigurationTest { + + private static final String TOKEN_URI = "https://auth.example.com/token"; + private static final String CLIENT_ID = "test-client-id"; + private static final String CLIENT_SECRET = "test-client-secret"; + private static final String SCOPE = "read write"; + private static final String SUBJECT_TOKEN = "subject-token-value"; + public static final String RESOURCE = "https://api.example.com/resource"; + + @FunctionalInterface + interface BuilderConfigurer { + void configure(OAuthConfiguration.Builder builder) throws SQLException; + } + + static Stream createFlowCases() { + return Stream.of( + Arguments.of( + Named.of( + "string flow", (BuilderConfigurer) builder -> builder.flow("client_credentials"))), + Arguments.of( + Named.of( + "uppercase string flow", + (BuilderConfigurer) builder -> builder.flow("CLIENT_CREDENTIALS")))); + } + + @ParameterizedTest + @MethodSource("createFlowCases") + public void testCreateFlowConfiguration(BuilderConfigurer flowConfigurer) throws SQLException { + OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder(); + flowConfigurer.configure(builder); + OAuthConfiguration config = + builder.tokenUri(TOKEN_URI).clientId(CLIENT_ID).clientSecret(CLIENT_SECRET).build(); + + // Verify configuration creates correct provider type + OAuthTokenProvider provider = config.createTokenProvider(); + assertInstanceOf(ClientCredentialsTokenProvider.class, provider); + } + + @Test + public void testCreateClientCredentialsTokenProvider() throws SQLException { + OAuthConfiguration config = + new OAuthConfiguration.Builder() + .flow("client_credentials") + .tokenUri(TOKEN_URI) + .clientId(CLIENT_ID) + .clientSecret(CLIENT_SECRET) + .scope(SCOPE) + .build(); + + OAuthTokenProvider provider = config.createTokenProvider(); + + assertNotNull(provider); + assertInstanceOf(ClientCredentialsTokenProvider.class, provider); + + ClientCredentialsTokenProvider ccProvider = (ClientCredentialsTokenProvider) provider; + assertEquals(URI.create(TOKEN_URI), ccProvider.tokenUri); + assertEquals(CLIENT_ID, ccProvider.clientAuth.getClientID().getValue()); + assertEquals(Scope.parse(SCOPE), ccProvider.scope); + } + + @Test + public void testCreateTokenExchangeTokenProviderWithAllOptions() throws SQLException { + String subjectTokenType = "urn:ietf:params:oauth:token-type:access_token"; + String actorToken = "actor-token-value"; + String actorTokenType = "urn:ietf:params:oauth:token-type:jwt"; + String audience = "https://api.example.com"; + String requestedTokenType = "urn:ietf:params:oauth:token-type:access_token"; + + OAuthConfiguration config = + new OAuthConfiguration.Builder() + .flow("token_exchange") + .tokenUri(TOKEN_URI) + .scope(SCOPE) + .clientId(CLIENT_ID) + .clientSecret(CLIENT_SECRET) + .resource(RESOURCE) + .subjectToken(SUBJECT_TOKEN) + .subjectTokenType(subjectTokenType) + .actorToken(actorToken) + .actorTokenType(actorTokenType) + .audience(audience) + .requestedTokenType(requestedTokenType) + .build(); + + OAuthTokenProvider provider = config.createTokenProvider(); + + assertNotNull(provider); + assertInstanceOf(TokenExchangeTokenProvider.class, provider); + + TokenExchangeTokenProvider teProvider = (TokenExchangeTokenProvider) provider; + assertEquals(URI.create(TOKEN_URI), teProvider.tokenUri); + assertNotNull(teProvider.grant); + assertEquals(SUBJECT_TOKEN, teProvider.grant.getSubjectToken().getValue()); + assertEquals(subjectTokenType, teProvider.grant.getSubjectTokenType().getURI().toString()); + assertEquals(actorToken, teProvider.grant.getActorToken().getValue()); + assertEquals(actorTokenType, teProvider.grant.getActorTokenType().getURI().toString()); + assertNotNull(teProvider.grant.getAudience()); + assertEquals(1, teProvider.grant.getAudience().size()); + assertEquals(audience, teProvider.grant.getAudience().get(0).getValue()); + assertEquals(requestedTokenType, teProvider.grant.getRequestedTokenType().getURI().toString()); + assertEquals(Scope.parse(SCOPE), teProvider.scope); + assertEquals(Collections.singletonList(URI.create(RESOURCE)), teProvider.resources); + + assertEquals(CLIENT_ID, teProvider.clientAuth.getClientID().getValue()); + } + + static Stream generalValidationErrorCases() { + return Stream.of( + Arguments.of( + Named.of( + "null flow", + (BuilderConfigurer) builder -> builder.flow((String) null).tokenUri(TOKEN_URI)), + "OAuth flow cannot be null or empty"), + Arguments.of( + Named.of( + "empty flow", (BuilderConfigurer) builder -> builder.flow("").tokenUri(TOKEN_URI)), + "OAuth flow cannot be null or empty"), + Arguments.of( + Named.of( + "invalid flow", + (BuilderConfigurer) builder -> builder.flow("invalid_flow").tokenUri(TOKEN_URI)), + "Unsupported OAuth flow: invalid_flow"), + Arguments.of( + Named.of( + "null tokenUri", + (BuilderConfigurer) + builder -> + builder + .flow("client_credentials") + .tokenUri((String) null) + .clientId(CLIENT_ID) + .clientSecret(CLIENT_SECRET)), + "Token URI cannot be null or empty"), + Arguments.of( + Named.of( + "empty tokenUri", + (BuilderConfigurer) + builder -> + builder + .flow("client_credentials") + .tokenUri("") + .clientId(CLIENT_ID) + .clientSecret(CLIENT_SECRET)), + "Token URI cannot be null or empty"), + Arguments.of( + Named.of( + "invalid tokenUri", + (BuilderConfigurer) + builder -> + builder + .flow("client_credentials") + .tokenUri("not a valid uri ://") + .clientId(CLIENT_ID) + .clientSecret(CLIENT_SECRET)), + null), + Arguments.of( + Named.of( + "invalid tokenUri", + (BuilderConfigurer) + builder -> + builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)), + // null means verify exception has message and cause + "clientSecret is required for client_credentials flow")); + } + + @ParameterizedTest + @MethodSource("generalValidationErrorCases") + public void testGeneralValidationErrors(BuilderConfigurer configurer, String expectedMessage) { + SQLException exception = + assertThrows( + SQLException.class, + () -> { + OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder(); + configurer.configure(builder); + builder.build(); + }); + + if (expectedMessage != null) { + assertEquals(expectedMessage, exception.getMessage()); + } else { + assertNotNull(exception.getMessage()); + assertNotNull(exception.getCause()); + } + } + + static Stream flowSpecificValidationErrorCases() { + return Stream.of( + // client_credentials flow validation + Arguments.of( + Named.of( + "client_credentials: missing clientId", + (BuilderConfigurer) + builder -> + builder + .flow("client_credentials") + .tokenUri(TOKEN_URI) + .clientSecret(CLIENT_SECRET)), + "clientId is required for client_credentials flow"), + Arguments.of( + Named.of( + "client_credentials: missing clientSecret", + (BuilderConfigurer) + builder -> + builder.flow("client_credentials").tokenUri(TOKEN_URI).clientId(CLIENT_ID)), + "clientSecret is required for client_credentials flow"), + // token_exchange flow validation + Arguments.of( + Named.of( + "token_exchange: missing subjectToken", + (BuilderConfigurer) builder -> builder.flow("token_exchange").tokenUri(TOKEN_URI)), + "subjectToken is required for token_exchange flow"), + Arguments.of( + Named.of( + "token_exchange: empty subjectToken", + (BuilderConfigurer) + builder -> + builder + .flow("token_exchange") + .tokenUri(TOKEN_URI) + .subjectToken("") + .subjectTokenType("urn:ietf:params:oauth:token-type:access_token")), + "subjectToken is required for token_exchange flow"), + Arguments.of( + Named.of( + "token_exchange: missing subjectTokenType", + (BuilderConfigurer) + builder -> + builder + .flow("token_exchange") + .tokenUri(TOKEN_URI) + .subjectToken(SUBJECT_TOKEN)), + "subjectTokenType is required for token_exchange flow"), + Arguments.of( + Named.of( + "token_exchange: empty subjectTokenType", + (BuilderConfigurer) + builder -> + builder + .flow("token_exchange") + .tokenUri(TOKEN_URI) + .subjectToken(SUBJECT_TOKEN) + .subjectTokenType("")), + "subjectTokenType is required for token_exchange flow")); + } + + @ParameterizedTest + @MethodSource("flowSpecificValidationErrorCases") + public void testFlowSpecificValidationErrors( + BuilderConfigurer configurer, String expectedMessage) { + SQLException exception = + assertThrows( + SQLException.class, + () -> { + OAuthConfiguration.Builder builder = new OAuthConfiguration.Builder(); + configurer.configure(builder); + builder.build(); + }); + + assertEquals(expectedMessage, exception.getMessage()); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java new file mode 100644 index 0000000000..1a33f7f0ae --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthCredentialWriterTest.java @@ -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.arrow.driver.jdbc.client.oauth; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.SQLException; +import org.apache.arrow.flight.CallHeaders; +import org.apache.arrow.flight.FlightCallHeaders; +import org.apache.arrow.flight.auth2.Auth2Constants; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** Tests for {@link OAuthCredentialWriter}. */ +@ExtendWith(MockitoExtension.class) +public class OAuthCredentialWriterTest { + + @Mock private OAuthTokenProvider mockTokenProvider; + + @Test + public void testConstructorRejectsNullTokenProvider() { + assertThrows(NullPointerException.class, () -> new OAuthCredentialWriter(null)); + } + + @Test + public void testAcceptWritesBearerTokenToHeaders() throws SQLException { + String testToken = "test-access-token-12345"; + when(mockTokenProvider.getValidToken()).thenReturn(testToken); + + OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider); + CallHeaders headers = new FlightCallHeaders(); + + writer.accept(headers); + + verify(mockTokenProvider).getValidToken(); + assertEquals( + Auth2Constants.BEARER_PREFIX + testToken, headers.get(Auth2Constants.AUTHORIZATION_HEADER)); + } + + @Test + public void testAcceptThrowsOAuthTokenExceptionOnSQLException() throws SQLException { + SQLException sqlException = new SQLException("Token fetch failed"); + when(mockTokenProvider.getValidToken()).thenThrow(sqlException); + + OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider); + CallHeaders headers = new FlightCallHeaders(); + + OAuthTokenException exception = + assertThrows(OAuthTokenException.class, () -> writer.accept(headers)); + + assertEquals("Failed to obtain OAuth token", exception.getMessage()); + assertEquals(sqlException, exception.getCause()); + } + + @Test + public void testAcceptCallsTokenProviderEachTime() throws SQLException { + when(mockTokenProvider.getValidToken()) + .thenReturn("token1") + .thenReturn("token2") + .thenReturn("token3"); + + OAuthCredentialWriter writer = new OAuthCredentialWriter(mockTokenProvider); + + CallHeaders headers1 = new FlightCallHeaders(); + writer.accept(headers1); + assertEquals("Bearer token1", headers1.get(Auth2Constants.AUTHORIZATION_HEADER)); + + CallHeaders headers2 = new FlightCallHeaders(); + writer.accept(headers2); + assertEquals("Bearer token2", headers2.get(Auth2Constants.AUTHORIZATION_HEADER)); + + CallHeaders headers3 = new FlightCallHeaders(); + writer.accept(headers3); + assertEquals("Bearer token3", headers3.get(Auth2Constants.AUTHORIZATION_HEADER)); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCacheTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCacheTest.java new file mode 100644 index 0000000000..8e818967a5 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCacheTest.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.client.utils; + +import static org.junit.jupiter.api.Assertions.*; + +import org.apache.arrow.flight.Location; +import org.junit.jupiter.api.Test; + +class FlightClientCacheTest { + @Test + void basicOperation() { + FlightClientCache cache = new FlightClientCache(); + + Location location1 = Location.forGrpcInsecure("localhost", 8080); + Location location2 = Location.forGrpcInsecure("localhost", 8081); + + assertFalse(cache.isDud(location1.toString())); + assertFalse(cache.isDud(location2.toString())); + + cache.markLocationAsReachable(location1.toString()); + assertFalse(cache.isDud(location1.toString())); + assertFalse(cache.isDud(location2.toString())); + + cache.markLocationAsDud(location1.toString()); + assertTrue(cache.isDud(location1.toString())); + assertFalse(cache.isDud(location2.toString())); + + cache.markLocationAsDud(location2.toString()); + assertTrue(cache.isDud(location1.toString())); + assertTrue(cache.isDud(location2.toString())); + + cache.markLocationAsReachable(location1.toString()); + assertFalse(cache.isDud(location1.toString())); + assertTrue(cache.isDud(location2.toString())); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueueTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueueTest.java new file mode 100644 index 0000000000..0603f86e59 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueueTest.java @@ -0,0 +1,63 @@ +/* + * 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.arrow.driver.jdbc.client.utils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.flight.Location; +import org.junit.jupiter.api.Test; + +class FlightLocationQueueTest { + @Test + void basicOperation() { + Location location1 = Location.forGrpcInsecure("localhost", 8080); + Location location2 = Location.forGrpcInsecure("localhost", 8081); + FlightLocationQueue queue = new FlightLocationQueue(null, List.of(location1, location2)); + assertTrue(queue.hasNext()); + assertEquals(location1, queue.next()); + assertTrue(queue.hasNext()); + assertEquals(location2, queue.next()); + assertFalse(queue.hasNext()); + } + + @Test + void badAfterGood() { + Location location1 = Location.forGrpcInsecure("localhost", 8080); + Location location2 = Location.forGrpcInsecure("localhost", 8081); + FlightClientCache cache = new FlightClientCache(); + cache.markLocationAsDud(location1.toString()); + FlightLocationQueue queue = new FlightLocationQueue(cache, List.of(location1, location2)); + assertTrue(queue.hasNext()); + assertEquals(location2, queue.next()); + assertTrue(queue.hasNext()); + assertEquals(location1, queue.next()); + assertFalse(queue.hasNext()); + } + + @Test + void iteratorInvariants() { + FlightLocationQueue empty = new FlightLocationQueue(null, Collections.emptyList()); + assertFalse(empty.hasNext()); + assertThrows(NoSuchElementException.class, empty::next); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java new file mode 100644 index 0000000000..07751f0abc --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/converter/impl/UuidAvaticaParameterConverterTest.java @@ -0,0 +1,160 @@ +/* + * 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.arrow.driver.jdbc.converter.impl; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.sql.Types; +import java.util.UUID; +import org.apache.arrow.driver.jdbc.utils.RootAllocatorTestExtension; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.UuidUtility; +import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.ColumnMetaData; +import org.apache.calcite.avatica.remote.TypedValue; +import org.apache.calcite.avatica.util.ByteString; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Tests for {@link UuidAvaticaParameterConverter}. + * + *

Verifies that the converter correctly handles UUID parameter binding from JDBC to Arrow's UUID + * extension type. + */ +public class UuidAvaticaParameterConverterTest { + + @RegisterExtension + public static RootAllocatorTestExtension rootAllocatorTestExtension = + new RootAllocatorTestExtension(); + + private static final UUID TEST_UUID = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + + private UuidVector vector; + private UuidAvaticaParameterConverter converter; + + @BeforeEach + public void setUp() { + vector = new UuidVector("uuid_param", rootAllocatorTestExtension.getRootAllocator()); + vector.allocateNew(5); + converter = new UuidAvaticaParameterConverter(); + } + + @AfterEach + public void tearDown() { + vector.close(); + } + + @Test + public void testBindParameterWithUuidObject() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, TEST_UUID); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithUuidString() { + String uuidString = "550e8400-e29b-41d4-a716-446655440000"; + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.STRING, uuidString); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithByteArray() { + byte[] uuidBytes = UuidUtility.getBytesFromUUID(TEST_UUID); + ByteString byteString = new ByteString(uuidBytes); + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertThat(vector.getObject(0), is(TEST_UUID)); + } + + @Test + public void testBindParameterWithNullValue() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, null); + + boolean result = converter.bindParameter(vector, typedValue, 0); + + assertTrue(result); + assertTrue(vector.isNull(0)); + assertThat(vector.getObject(0), nullValue()); + } + + @Test + public void testBindParameterWithInvalidByteArrayLength() { + byte[] invalidBytes = new byte[8]; // Should be 16 bytes + ByteString byteString = new ByteString(invalidBytes); + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.BYTE_STRING, byteString); + + assertThrows( + IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0)); + } + + @Test + public void testBindParameterWithInvalidType() { + TypedValue typedValue = TypedValue.ofLocal(ColumnMetaData.Rep.INTEGER, 12345); + + assertThrows( + IllegalArgumentException.class, () -> converter.bindParameter(vector, typedValue, 0)); + } + + @Test + public void testBindParameterMultipleValues() { + UUID uuid1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + UUID uuid2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + UUID uuid3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid1), 0); + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid2), 1); + converter.bindParameter(vector, TypedValue.ofLocal(ColumnMetaData.Rep.OBJECT, uuid3), 2); + + assertThat(vector.getObject(0), is(uuid1)); + assertThat(vector.getObject(1), is(uuid2)); + assertThat(vector.getObject(2), is(uuid3)); + } + + @Test + public void testCreateParameter() { + Field uuidField = new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null); + + AvaticaParameter parameter = converter.createParameter(uuidField); + + assertThat(parameter.name, is("uuid_col")); + assertThat(parameter.parameterType, is(Types.OTHER)); + assertThat(parameter.typeName, is("OTHER")); + assertThat(parameter.className, equalTo(UUID.class.getCanonicalName())); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImplTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImplTest.java index c780d53fab..ecce7708c0 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImplTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImplTest.java @@ -24,6 +24,7 @@ import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.PORT; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.THREAD_POOL_SIZE; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.USER; +import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.USE_CLIENT_CACHE; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.USE_ENCRYPTION; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; @@ -122,6 +123,12 @@ public static Stream provideParameters() { 5000, Duration.ofMillis(5000), (Function) - ArrowFlightConnectionConfigImpl::getConnectTimeout)); + ArrowFlightConnectionConfigImpl::getConnectTimeout), + Arguments.of( + USE_CLIENT_CACHE, + false, + false, + (Function) + ArrowFlightConnectionConfigImpl::useClientCache)); } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java index f6f549b5ed..f128ca7c73 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/ConvertUtilsTest.java @@ -46,6 +46,7 @@ public void testShouldSetOnColumnMetaDataBuilder() { .isSearchable(true) .precision(20) .scale(10) + .remarks("test column") .build(); ConvertUtils.setOnColumnMetaDataBuilder(builder, expectedColumnMetaData.getMetadataMap()); assertBuilder(builder, expectedColumnMetaData); @@ -68,6 +69,19 @@ public void testShouldConvertArrowFieldsToColumnMetaDataList() { .tableName("table1") .build() .getMetadataMap()), + null), + new Field( + "col2", + new FieldType( + true, + ArrowType.Utf8View.INSTANCE, + null, + new FlightSqlColumnMetadata.Builder() + .catalogName("catalog1") + .schemaName("schema1") + .tableName("table1") + .build() + .getMetadataMap()), null)); final List expectedColumnMetaData = @@ -77,6 +91,25 @@ public void testShouldConvertArrowFieldsToColumnMetaDataList() { .setCatalogName("catalog1") .setSchemaName("schema1") .setTableName("table1") + .setColumnName("col1") + .setType( + Common.AvaticaType.newBuilder() + .setId(SqlTypes.getSqlTypeIdFromArrowType(ArrowType.Utf8.INSTANCE)) + .setName(SqlTypes.getSqlTypeNameFromArrowType(ArrowType.Utf8.INSTANCE)) + .build()) + .build()), + ColumnMetaData.fromProto( + Common.ColumnMetaData.newBuilder() + .setCatalogName("catalog1") + .setSchemaName("schema1") + .setTableName("table1") + .setColumnName("col2") + .setType( + Common.AvaticaType.newBuilder() + .setId(SqlTypes.getSqlTypeIdFromArrowType(ArrowType.Utf8View.INSTANCE)) + .setName( + SqlTypes.getSqlTypeNameFromArrowType(ArrowType.Utf8View.INSTANCE)) + .build()) .build())); final List actualColumnMetaData = @@ -94,6 +127,8 @@ private void assertColumnMetaData( assertThat(expectedColumnMetaData.catalogName, equalTo(actualColumnMetaData.catalogName)); assertThat(expectedColumnMetaData.schemaName, equalTo(actualColumnMetaData.schemaName)); assertThat(expectedColumnMetaData.tableName, equalTo(actualColumnMetaData.tableName)); + assertThat(expectedColumnMetaData.columnName, equalTo(actualColumnMetaData.columnName)); + assertThat(expectedColumnMetaData.type, equalTo(actualColumnMetaData.type)); assertThat(expectedColumnMetaData.readOnly, equalTo(actualColumnMetaData.readOnly)); assertThat(expectedColumnMetaData.autoIncrement, equalTo(actualColumnMetaData.autoIncrement)); assertThat(expectedColumnMetaData.precision, equalTo(actualColumnMetaData.precision)); @@ -119,5 +154,6 @@ private void assertBuilder( assertThat(flightSqlColumnMetaData.isReadOnly(), equalTo(builder.getReadOnly())); assertThat(precision == null ? 0 : precision, equalTo(builder.getPrecision())); assertThat(scale == null ? 0 : scale, equalTo(builder.getScale())); + assertThat(flightSqlColumnMetaData.getRemarks(), equalTo(builder.getLabel())); } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java index 8197d7d95f..7c17755693 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/CoreMockedSqlProducers.java @@ -28,8 +28,10 @@ import java.sql.SQLException; import java.sql.Timestamp; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.UUID; import java.util.function.Consumer; import java.util.stream.IntStream; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; @@ -40,10 +42,13 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.TimeStampMilliVector; import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; @@ -52,6 +57,7 @@ import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; +import org.apache.arrow.vector.util.UuidUtility; /** Standard {@link MockFlightSqlProducer} instances for tests. */ // TODO Remove this once all tests are refactor to use only the queries they need. @@ -62,6 +68,22 @@ public final class CoreMockedSqlProducers { public static final String LEGACY_CANCELLATION_SQL_CMD = "SELECT * FROM TAKES_FOREVER"; public static final String LEGACY_REGULAR_WITH_EMPTY_SQL_CMD = "SELECT * FROM TEST_EMPTIES"; + public static final String UUID_SQL_CMD = "SELECT * FROM UUID_TABLE"; + public static final String UUID_PREPARED_SELECT_SQL_CMD = + "SELECT * FROM UUID_TABLE WHERE uuid_col = ?"; + public static final String UUID_PREPARED_UPDATE_SQL_CMD = + "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?"; + + public static final UUID UUID_1 = UUID.fromString("550e8400-e29b-41d4-a716-446655440000"); + public static final UUID UUID_2 = UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"); + public static final UUID UUID_3 = UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"); + + public static final Schema UUID_SCHEMA = + new Schema( + ImmutableList.of( + new Field("id", new FieldType(true, new ArrowType.Int(32, true), null), null), + new Field("uuid_col", new FieldType(true, UuidType.INSTANCE, null), null))); + private CoreMockedSqlProducers() { // Prevent instantiation. } @@ -78,9 +100,109 @@ public static MockFlightSqlProducer getLegacyProducer() { addLegacyMetadataSqlCmdSupport(producer); addLegacyCancellationSqlCmdSupport(producer); addQueryWithEmbeddedEmptyRoot(producer); + addUuidSqlCmdSupport(producer); + addUuidPreparedSelectSqlCmdSupport(producer); + addUuidPreparedUpdateSqlCmdSupport(producer); return producer; } + /** + * Gets a {@link MockFlightSqlProducer} configured with UUID test data. + * + * @return a new producer with UUID support. + */ + public static MockFlightSqlProducer getUuidProducer() { + final MockFlightSqlProducer producer = new MockFlightSqlProducer(); + addUuidSqlCmdSupport(producer); + return producer; + } + + private static void addUuidPreparedUpdateSqlCmdSupport(final MockFlightSqlProducer producer) { + final String query = "UPDATE UUID_TABLE SET uuid_col = ? WHERE id = ?"; + final Schema parameterSchema = + new Schema( + Arrays.asList( + new Field("", new FieldType(true, UuidType.INSTANCE, null), null), + Field.nullable("", new ArrowType.Int(32, true)))); + + producer.addUpdateQuery(query, 1); + producer.addExpectedParameters( + UUID_PREPARED_UPDATE_SQL_CMD, + parameterSchema, + Collections.singletonList(Arrays.asList(CoreMockedSqlProducers.UUID_3, 1))); + } + + private static void addUuidPreparedSelectSqlCmdSupport(final MockFlightSqlProducer producer) { + final Schema parameterSchema = + new Schema( + Collections.singletonList( + new Field("", new FieldType(true, UuidType.INSTANCE, null), null))); + + final Consumer uuidResultProvider = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) { + root.allocateNew(); + IntVector idVector = (IntVector) root.getVector("id"); + UuidVector uuidVector = (UuidVector) root.getVector("uuid_col"); + idVector.setSafe(0, 1); + uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(CoreMockedSqlProducers.UUID_1)); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + }; + + producer.addSelectQuery( + UUID_PREPARED_SELECT_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider)); + producer.addExpectedParameters( + UUID_PREPARED_SELECT_SQL_CMD, + parameterSchema, + Collections.singletonList(Collections.singletonList(CoreMockedSqlProducers.UUID_1))); + } + + private static void addUuidSqlCmdSupport(final MockFlightSqlProducer producer) { + final Consumer uuidResultProvider = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(UUID_SCHEMA, allocator)) { + root.allocateNew(); + + IntVector idVector = (IntVector) root.getVector("id"); + UuidVector uuidVector = (UuidVector) root.getVector("uuid_col"); + + // Row 0: id=1, uuid=UUID_1 + idVector.setSafe(0, 1); + uuidVector.setSafe(0, UuidUtility.getBytesFromUUID(UUID_1)); + + // Row 1: id=2, uuid=UUID_2 + idVector.setSafe(1, 2); + uuidVector.setSafe(1, UuidUtility.getBytesFromUUID(UUID_2)); + + // Row 2: id=3, uuid=UUID_3 + idVector.setSafe(2, 3); + uuidVector.setSafe(2, UuidUtility.getBytesFromUUID(UUID_3)); + + // Row 3: id=4, uuid=NULL + idVector.setSafe(3, 4); + uuidVector.setNull(3); + + root.setRowCount(4); + listener.start(root); + listener.putNext(); + } finally { + listener.completed(); + } + }; + + producer.addSelectQuery( + UUID_SQL_CMD, UUID_SCHEMA, Collections.singletonList(uuidResultProvider)); + } + private static void addQueryWithEmbeddedEmptyRoot(final MockFlightSqlProducer producer) { final Schema querySchema = new Schema( diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java index a8874c4869..6627d91ab6 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java @@ -52,6 +52,9 @@ import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.SetSessionOptionsRequest; +import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.flight.sql.SqlInfoBuilder; @@ -84,6 +87,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.calcite.avatica.Meta.StatementType; +import org.checkerframework.checker.nullness.qual.Nullable; /** An ad-hoc {@link FlightSqlProducer} for tests. */ public final class MockFlightSqlProducer implements FlightSqlProducer { @@ -98,6 +102,7 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder(); private final Map parameterSchemas = new HashMap<>(); private final Map>> expectedParameterValues = new HashMap<>(); + private final Map isUpdateMap = new HashMap<>(); private final Map actionTypeCounter = new HashMap<>(); @@ -173,6 +178,40 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) { }); } + /** + * Registers a new {@link StatementType#SELECT} SQL query, optionally setting the is_update field. + * + * @param sqlCommand the SQL command under which to register the new query. + * @param schema the schema to use for the query result. + * @param resultProviders the result provider for this query. + * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset. + */ + public void addSelectQuery( + final String sqlCommand, + final Schema schema, + final List> resultProviders, + final @Nullable Boolean isUpdate) { + addSelectQuery(sqlCommand, schema, resultProviders); + if (isUpdate != null) { + isUpdateMap.put(sqlCommand, isUpdate); + } + } + + /** + * Registers a new {@link StatementType#UPDATE} SQL query, optionally setting the is_update field. + * + * @param sqlCommand the SQL command. + * @param updatedRows the number of rows affected. + * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset. + */ + public void addUpdateQuery( + final String sqlCommand, final long updatedRows, final @Nullable Boolean isUpdate) { + addUpdateQuery(sqlCommand, updatedRows); + if (isUpdate != null) { + isUpdateMap.put(sqlCommand, isUpdate); + } + } + /** * Adds a catalog query to the results. * @@ -244,6 +283,12 @@ public void createPreparedStatement( resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray())); } + // Set is_update field if present + final Boolean isUpdate = isUpdateMap.get(query); + if (isUpdate != null) { + resultBuilder.setIsUpdate(isUpdate); + } + listener.onNext(new Result(pack(resultBuilder.build()).toByteArray())); } catch (final Throwable t) { listener.onError(t); @@ -664,6 +709,22 @@ public SqlInfoBuilder getSqlInfoBuilder() { return sqlInfoBuilder; } + private final Map sessionOptions = new HashMap<>(); + + @Override + public void setSessionOptions( + final SetSessionOptionsRequest request, + final CallContext context, + final StreamListener listener) { + sessionOptions.putAll(request.getSessionOptions()); + listener.onNext(new SetSessionOptionsResult(Collections.emptyMap())); + listener.onCompleted(); + } + + public Map getSessionOptions() { + return sessionOptions; + } + private static final class TicketConversionUtils { private TicketConversionUtils() { // Prevent instantiation. diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java index 347e92a16c..4b299d63e0 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/RootAllocatorTestExtension.java @@ -19,6 +19,7 @@ import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.Random; +import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.stream.IntStream; import org.apache.arrow.memory.BufferAllocator; @@ -53,6 +54,7 @@ import org.apache.arrow.vector.UInt2Vector; import org.apache.arrow.vector.UInt4Vector; import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; @@ -60,6 +62,7 @@ import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.extension.AfterAllCallback; import org.junit.jupiter.api.extension.BeforeAllCallback; import org.junit.jupiter.api.extension.ExtensionContext; @@ -811,4 +814,23 @@ public FixedSizeListVector createFixedSizeListVector() { return valueVector; } + + /** + * Create a UuidVector to be used in the accessor tests. + * + * @return UuidVector + */ + public UuidVector createUuidVector() { + UuidVector valueVector = new UuidVector("", this.getRootAllocator()); + valueVector.allocateNew(3); + valueVector.setSafe( + 0, UuidUtility.getBytesFromUUID(UUID.fromString("550e8400-e29b-41d4-a716-446655440000"))); + valueVector.setSafe( + 1, UuidUtility.getBytesFromUUID(UUID.fromString("6ba7b810-9dad-11d1-80b4-00c04fd430c8"))); + valueVector.setSafe( + 2, UuidUtility.getBytesFromUUID(UUID.fromString("f47ac10b-58cc-4372-a567-0e02b2c3d479"))); + valueVector.setValueCount(3); + + return valueVector; + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java index a6dd6b3275..c4858d787d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/SqlTypesTest.java @@ -21,6 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import java.sql.Types; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.IntervalUnit; @@ -40,9 +41,11 @@ public void testGetSqlTypeIdFromArrowType() { assertEquals(Types.BINARY, getSqlTypeIdFromArrowType(new ArrowType.FixedSizeBinary(1024))); assertEquals(Types.VARBINARY, getSqlTypeIdFromArrowType(new ArrowType.Binary())); + assertEquals(Types.VARBINARY, getSqlTypeIdFromArrowType(new ArrowType.BinaryView())); assertEquals(Types.LONGVARBINARY, getSqlTypeIdFromArrowType(new ArrowType.LargeBinary())); assertEquals(Types.VARCHAR, getSqlTypeIdFromArrowType(new ArrowType.Utf8())); + assertEquals(Types.VARCHAR, getSqlTypeIdFromArrowType(new ArrowType.Utf8View())); assertEquals(Types.LONGVARCHAR, getSqlTypeIdFromArrowType(new ArrowType.LargeUtf8())); assertEquals(Types.DATE, getSqlTypeIdFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND))); @@ -83,6 +86,8 @@ public void testGetSqlTypeIdFromArrowType() { assertEquals(Types.JAVA_OBJECT, getSqlTypeIdFromArrowType(new ArrowType.Map(true))); assertEquals(Types.NULL, getSqlTypeIdFromArrowType(new ArrowType.Null())); + + assertEquals(Types.OTHER, getSqlTypeIdFromArrowType(UuidType.INSTANCE)); } @Test @@ -94,9 +99,11 @@ public void testGetSqlTypeNameFromArrowType() { assertEquals("BINARY", getSqlTypeNameFromArrowType(new ArrowType.FixedSizeBinary(1024))); assertEquals("VARBINARY", getSqlTypeNameFromArrowType(new ArrowType.Binary())); + assertEquals("VARBINARY", getSqlTypeNameFromArrowType(new ArrowType.BinaryView())); assertEquals("LONGVARBINARY", getSqlTypeNameFromArrowType(new ArrowType.LargeBinary())); assertEquals("VARCHAR", getSqlTypeNameFromArrowType(new ArrowType.Utf8())); + assertEquals("VARCHAR", getSqlTypeNameFromArrowType(new ArrowType.Utf8View())); assertEquals("LONGVARCHAR", getSqlTypeNameFromArrowType(new ArrowType.LargeUtf8())); assertEquals("DATE", getSqlTypeNameFromArrowType(new ArrowType.Date(DateUnit.MILLISECOND))); @@ -136,5 +143,7 @@ public void testGetSqlTypeNameFromArrowType() { assertEquals("JAVA_OBJECT", getSqlTypeNameFromArrowType(new ArrowType.Map(true))); assertEquals("NULL", getSqlTypeNameFromArrowType(new ArrowType.Null())); + + assertEquals("OTHER", getSqlTypeNameFromArrowType(UuidType.INSTANCE)); } } diff --git a/flight/flight-sql-jdbc-driver/pom.xml b/flight/flight-sql-jdbc-driver/pom.xml index 559c42597d..801089d090 100644 --- a/flight/flight-sql-jdbc-driver/pom.xml +++ b/flight/flight-sql-jdbc-driver/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-flight - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT flight-sql-jdbc-driver @@ -138,7 +138,7 @@ under the License. - + META-INF/LICENSE.txt src/shade/LICENSE.txt @@ -159,6 +159,7 @@ under the License. org.apache.calcite.avatica:* META-INF/services/java.sql.Driver + META-INF/README.txt diff --git a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt index 1048fa8a8d..8476bd9995 100644 --- a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt +++ b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt @@ -227,14 +227,14 @@ License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- -This binary artifact contains Guava 33.4.5-jre. +This binary artifact contains Guava 33.4.8-jre. Home page: https://github.com/google/guava License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- -This binary artifact contains Protobuf 4.30.1. +This binary artifact contains Protobuf 4.30.2. Copyright: Copyright 2008 Google Inc. All rights reserved. Home page: https://protobuf.dev/ @@ -345,6 +345,14 @@ License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- +This binary artifact contains Nimbus OAuth 2.0 SDK with OpenID Connect extensions 11.20.1. + +Copyright: Copyright 2012-2024 Connect2id Ltd. +Home page: https://connect2id.com/products/nimbus-oauth-openid-connect-sdk +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + This binary artifact contains Bouncycastle 1.80. Copyright: Copyright (c) 2000-2024 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org). diff --git a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java index c1bd111fb9..145744ad38 100644 --- a/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java +++ b/flight/flight-sql-jdbc-driver/src/test/java/org/apache/arrow/driver/jdbc/ITDriverJarValidation.java @@ -70,6 +70,10 @@ public class ITDriverJarValidation { "LICENSE.txt", "NOTICE.txt", "arrow-git.properties", + "iso3166_1alpha2-codes.properties", + "iso3166_1alpha3-codes.properties", + "iso3166_1alpha-2-3-map.properties", + "iso3166_3-codes.properties", "properties/flight.properties", "META-INF/io.netty.versions.properties", "META-INF/MANIFEST.MF", diff --git a/flight/flight-sql/pom.xml b/flight/flight-sql/pom.xml index 5d8e720fd6..b7c8931391 100644 --- a/flight/flight-sql/pom.xml +++ b/flight/flight-sql/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-flight - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT flight-sql @@ -95,7 +95,7 @@ under the License. org.apache.commons commons-dbcp2 - 2.13.0 + 2.14.0 test @@ -107,25 +107,25 @@ under the License. org.apache.commons commons-pool2 - 2.12.1 + 2.13.1 test org.apache.commons commons-text - 1.13.0 - test - - - org.hamcrest - hamcrest + 1.15.0 test commons-cli commons-cli - 1.9.0 + 1.11.0 true + + org.assertj + assertj-core + test + diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java index 623f9311e8..0af09faee1 100644 --- a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java +++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java @@ -1284,6 +1284,19 @@ public Schema getParameterSchema() { return parameterSchema; } + /** + * Returns whether the server indicated this prepared statement is an update query. + * + * @return true if the server indicated this is an update query, false if the server indicated + * this is a select query, or null if the server did not provide this information. + */ + public Boolean isUpdate() { + if (preparedStatementResult.hasIsUpdate()) { + return preparedStatementResult.getIsUpdate(); + } + return null; + } + /** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */ public SchemaResult fetchSchema(CallOption... options) { checkOpen(); diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlColumnMetadata.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlColumnMetadata.java index 1bcc55a660..3a969e10cf 100644 --- a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlColumnMetadata.java +++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlColumnMetadata.java @@ -53,6 +53,7 @@ public class FlightSqlColumnMetadata { private static final String IS_CASE_SENSITIVE = "ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE"; private static final String IS_READ_ONLY = "ARROW:FLIGHT:SQL:IS_READ_ONLY"; private static final String IS_SEARCHABLE = "ARROW:FLIGHT:SQL:IS_SEARCHABLE"; + private static final String REMARKS = "ARROW:FLIGHT:SQL:REMARKS"; private static final String BOOLEAN_TRUE_STR = "1"; private static final String BOOLEAN_FALSE_STR = "0"; @@ -193,6 +194,15 @@ public Boolean isSearchable() { return stringToBoolean(value); } + /** + * Returns the comment describing the column. + * + * @return The comment describing the column. + */ + public String getRemarks() { + return metadataMap.get(REMARKS); + } + /** Builder of FlightSqlColumnMetadata, used on FlightSqlProducer implementations. */ public static class Builder { private final Map metadataMap; @@ -312,6 +322,17 @@ public Builder isSearchable(boolean isSearchable) { return this; } + /** + * Sets the comment describing the column. + * + * @param remarks The comment describing the column. + * @return This builder. + */ + public Builder remarks(String remarks) { + metadataMap.put(REMARKS, remarks); + return this; + } + /** * Builds a new instance of FlightSqlColumnMetadata. * diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java index 3f769363fb..e2934ab1e9 100644 --- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java +++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java @@ -21,10 +21,7 @@ import static java.util.Collections.singletonList; import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults; import static org.apache.arrow.util.AutoCloseables.close; -import static org.hamcrest.CoreMatchers.containsString; -import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.CoreMatchers.notNullValue; -import static org.hamcrest.CoreMatchers.nullValue; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -40,6 +37,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.stream.IntStream; import org.apache.arrow.flight.CancelFlightInfoRequest; @@ -76,8 +74,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.VectorBatchAppender; -import org.hamcrest.Matcher; -import org.hamcrest.MatcherAssert; +import org.assertj.core.api.Condition; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -247,16 +244,15 @@ private static List> getNonConformingResultsForGetSqlInfo( @Test public void testGetTablesSchema() { final FlightInfo info = sqlClient.getTables(null, null, null, null, true); - MatcherAssert.assertThat( - info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA)); } @Test public void testGetTablesSchemaExcludeSchema() { final FlightInfo info = sqlClient.getTables(null, null, null, null, false); - MatcherAssert.assertThat( - info.getSchemaOptional(), - is(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)); } @Test @@ -266,8 +262,8 @@ public void testGetTablesResultNoSchema() throws Exception { sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) { assertAll( () -> { - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)); + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA); }, () -> { final List> results = getResults(stream); @@ -301,7 +297,7 @@ public void testGetTablesResultNoSchema() throws Exception { asList(null /* TODO No catalog yet */, "SYSIBM", "SYSDUMMY1", "SYSTEM TABLE"), asList(null /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE"), asList(null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE")); - MatcherAssert.assertThat(results, is(expectedResults)); + assertThat(results).isEqualTo(expectedResults); }); } } @@ -318,8 +314,8 @@ public void testGetTablesResultFilteredNoSchema() throws Exception { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA), () -> { final List> results = getResults(stream); final List> expectedResults = @@ -327,7 +323,7 @@ public void testGetTablesResultFilteredNoSchema() throws Exception { // catalog_name | schema_name | table_name | table_type | table_schema asList(null /* TODO No catalog yet */, "APP", "FOREIGNTABLE", "TABLE"), asList(null /* TODO No catalog yet */, "APP", "INTTABLE", "TABLE")); - MatcherAssert.assertThat(results, is(expectedResults)); + assertThat(results).isEqualTo(expectedResults); }); } } @@ -343,11 +339,9 @@ public void testGetTablesResultFilteredWithSchema() throws Exception { .getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA)), + assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA), () -> { - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA)); + assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA); final List> results = getResults(stream); final List> expectedResults = ImmutableList.of( @@ -487,7 +481,7 @@ public void testGetTablesResultFilteredWithSchema() throws Exception { .getMetadataMap()), null))) .toJson())); - MatcherAssert.assertThat(results, is(expectedResults)); + assertThat(results).isEqualTo(expectedResults); }); } } @@ -498,11 +492,11 @@ public void testSimplePreparedStatementSchema() throws Exception { assertAll( () -> { final Schema actualSchema = preparedStatement.getResultSetSchema(); - MatcherAssert.assertThat(actualSchema, is(SCHEMA_INT_TABLE)); + assertThat(actualSchema).isEqualTo(SCHEMA_INT_TABLE); }, () -> { final FlightInfo info = preparedStatement.execute(); - MatcherAssert.assertThat(info.getSchemaOptional(), is(Optional.of(SCHEMA_INT_TABLE))); + assertThat(info.getSchemaOptional()).isEqualTo(Optional.of(SCHEMA_INT_TABLE)); }); } } @@ -513,10 +507,8 @@ public void testSimplePreparedStatementResults() throws Exception { final FlightStream stream = sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) { assertAll( - () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)), - () -> - MatcherAssert.assertThat( - getResults(stream), is(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY))); + () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), + () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY)); } } @@ -538,10 +530,8 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket()); assertAll( - () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)), - () -> - MatcherAssert.assertThat( - getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING))); + () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), + () -> assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)); } } } @@ -579,8 +569,8 @@ public void testSimplePreparedStatementUpdateResults() throws SQLException { deletedRows = deletePrepare.executeUpdate(); } assertAll( - () -> MatcherAssert.assertThat(updatedRows, is(10L)), - () -> MatcherAssert.assertThat(deletedRows, is(10L))); + () -> assertThat(updatedRows).isEqualTo(10L), + () -> assertThat(deletedRows).isEqualTo(10L)); } } } @@ -647,7 +637,7 @@ public void testBulkIngest() throws IOException { null, null)); - MatcherAssert.assertThat(updatedRows, is(-1L)); + assertThat(updatedRows).isEqualTo(-1L); // Ingest directly using VectorSchemaRoot populateNext10RowsInIngestRootBatch( @@ -672,7 +662,7 @@ public void testBulkIngest() throws IOException { deletedRows = deletePrepare.executeUpdate(); } - MatcherAssert.assertThat(deletedRows, is(30L)); + assertThat(deletedRows).isEqualTo(30L); } } } @@ -709,8 +699,7 @@ public void testSimplePreparedStatementUpdateResultsWithoutParameters() throws S final long deletedRows = deletePrepare.executeUpdate(); assertAll( - () -> MatcherAssert.assertThat(updatedRows, is(1L)), - () -> MatcherAssert.assertThat(deletedRows, is(1L))); + () -> assertThat(updatedRows).isEqualTo(1L), () -> assertThat(deletedRows).isEqualTo(1L)); } } @@ -719,19 +708,19 @@ public void testSimplePreparedStatementClosesProperly() { final PreparedStatement preparedStatement = sqlClient.prepare("SELECT * FROM intTable"); assertAll( () -> { - MatcherAssert.assertThat(preparedStatement.isClosed(), is(false)); + assertThat(preparedStatement.isClosed()).isEqualTo(false); }, () -> { preparedStatement.close(); - MatcherAssert.assertThat(preparedStatement.isClosed(), is(true)); + assertThat(preparedStatement.isClosed()).isEqualTo(true); }); } @Test public void testGetCatalogsSchema() { final FlightInfo info = sqlClient.getCatalogs(); - MatcherAssert.assertThat( - info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA)); } @Test @@ -740,11 +729,11 @@ public void testGetCatalogsResults() throws Exception { sqlClient.getStream(sqlClient.getCatalogs().getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA), () -> { List> catalogs = getResults(stream); - MatcherAssert.assertThat(catalogs, is(emptyList())); + assertThat(catalogs).isEqualTo(emptyList()); }); } } @@ -752,9 +741,8 @@ public void testGetCatalogsResults() throws Exception { @Test public void testGetTableTypesSchema() { final FlightInfo info = sqlClient.getTableTypes(); - MatcherAssert.assertThat( - info.getSchemaOptional(), - is(Optional.of(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)); } @Test @@ -763,8 +751,8 @@ public void testGetTableTypesResult() throws Exception { sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) { assertAll( () -> { - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)); + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA); }, () -> { final List> tableTypes = getResults(stream); @@ -775,7 +763,7 @@ public void testGetTableTypesResult() throws Exception { singletonList("SYSTEM TABLE"), singletonList("TABLE"), singletonList("VIEW")); - MatcherAssert.assertThat(tableTypes, is(expectedTableTypes)); + assertThat(tableTypes).isEqualTo(expectedTableTypes); }); } } @@ -783,8 +771,8 @@ public void testGetTableTypesResult() throws Exception { @Test public void testGetSchemasSchema() { final FlightInfo info = sqlClient.getSchemas(null, null); - MatcherAssert.assertThat( - info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA)); } @Test @@ -793,8 +781,7 @@ public void testGetSchemasResult() throws Exception { sqlClient.getStream(sqlClient.getSchemas(null, null).getEndpoints().get(0).getTicket())) { assertAll( () -> { - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA)); + assertThat(stream.getSchema()).isEqualTo(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA); }, () -> { final List> schemas = getResults(stream); @@ -812,7 +799,7 @@ public void testGetSchemasResult() throws Exception { asList(null /* TODO Add catalog. */, "SYSIBM"), asList(null /* TODO Add catalog. */, "SYSPROC"), asList(null /* TODO Add catalog. */, "SYSSTAT")); - MatcherAssert.assertThat(schemas, is(expectedSchemas)); + assertThat(schemas).isEqualTo(expectedSchemas); }); } } @@ -825,24 +812,24 @@ public void testGetPrimaryKey() { final List> results = getResults(stream); assertAll( - () -> MatcherAssert.assertThat(results.size(), is(1)), + () -> assertThat(results.size()).isEqualTo(1), () -> { final List result = results.get(0); assertAll( - () -> MatcherAssert.assertThat(result.get(0), is("")), - () -> MatcherAssert.assertThat(result.get(1), is("APP")), - () -> MatcherAssert.assertThat(result.get(2), is("INTTABLE")), - () -> MatcherAssert.assertThat(result.get(3), is("ID")), - () -> MatcherAssert.assertThat(result.get(4), is("1")), - () -> MatcherAssert.assertThat(result.get(5), notNullValue())); + () -> assertThat(result.get(0)).isEqualTo(""), + () -> assertThat(result.get(1)).isEqualTo("APP"), + () -> assertThat(result.get(2)).isEqualTo("INTTABLE"), + () -> assertThat(result.get(3)).isEqualTo("ID"), + () -> assertThat(result.get(4)).isEqualTo("1"), + () -> assertThat(result.get(5)).isNotNull()); }); } @Test public void testGetSqlInfoSchema() { final FlightInfo info = sqlClient.getSqlInfo(); - MatcherAssert.assertThat( - info.getSchemaOptional(), is(Optional.of(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA))); + assertThat(info.getSchemaOptional()) + .isEqualTo(Optional.of(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)); } @Test @@ -851,11 +838,11 @@ public void testGetSqlInfoResults() throws Exception { try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> - MatcherAssert.assertThat( - getNonConformingResultsForGetSqlInfo(getResults(stream)), is(emptyList()))); + assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream))) + .isEqualTo(emptyList())); } } @@ -866,11 +853,11 @@ public void testGetSqlInfoResultsWithSingleArg() throws Exception { try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> - MatcherAssert.assertThat( - getNonConformingResultsForGetSqlInfo(getResults(stream), arg), is(emptyList()))); + assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), arg)) + .isEqualTo(emptyList())); } } @@ -895,11 +882,11 @@ public void testGetSqlInfoResultsWithManyArgs() throws Exception { try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), () -> - MatcherAssert.assertThat( - getNonConformingResultsForGetSqlInfo(getResults(stream), args), is(emptyList()))); + assertThat(getNonConformingResultsForGetSqlInfo(getResults(stream), args)) + .isEqualTo(emptyList())); } } @@ -915,28 +902,30 @@ public void testGetCommandExportedKeys() throws Exception { final List> results = getResults(stream); - final List> matchers = + final List> matchers = asList( - nullValue(String.class), // pk_catalog_name - is("APP"), // pk_schema_name - is("FOREIGNTABLE"), // pk_table_name - is("ID"), // pk_column_name - nullValue(String.class), // fk_catalog_name - is("APP"), // fk_schema_name - is("INTTABLE"), // fk_table_name - is("FOREIGNID"), // fk_column_name - is("1"), // key_sequence - containsString("SQL"), // fk_key_name - containsString("SQL"), // pk_key_name - is("3"), // update_rule - is("3")); // delete_rule + new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"), + new Condition<>( + c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"), + new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"), + new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"), + new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"), + new Condition<>( + c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"), + new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"), + new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"), + new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"), + new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"), + new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3")); final List assertions = new ArrayList<>(); assertEquals(1, results.size()); for (int i = 0; i < matchers.size(); i++) { final String actual = results.get(0).get(i); - final Matcher expected = matchers.get(i); - assertions.add(() -> MatcherAssert.assertThat(actual, expected)); + final Condition expected = matchers.get(i); + assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } @@ -954,28 +943,30 @@ public void testGetCommandImportedKeys() throws Exception { final List> results = getResults(stream); - final List> matchers = + final List> matchers = asList( - nullValue(String.class), // pk_catalog_name - is("APP"), // pk_schema_name - is("FOREIGNTABLE"), // pk_table_name - is("ID"), // pk_column_name - nullValue(String.class), // fk_catalog_name - is("APP"), // fk_schema_name - is("INTTABLE"), // fk_table_name - is("FOREIGNID"), // fk_column_name - is("1"), // key_sequence - containsString("SQL"), // fk_key_name - containsString("SQL"), // pk_key_name - is("3"), // update_rule - is("3")); // delete_rule + new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"), + new Condition<>( + c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"), + new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"), + new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"), + new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"), + new Condition<>( + c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"), + new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"), + new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"), + new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"), + new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"), + new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3")); assertEquals(1, results.size()); final List assertions = new ArrayList<>(); for (int i = 0; i < matchers.size(); i++) { final String actual = results.get(0).get(i); - final Matcher expected = matchers.get(i); - assertions.add(() -> MatcherAssert.assertThat(actual, expected)); + final Condition expected = matchers.get(i); + assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } @@ -1431,7 +1422,7 @@ public void testGetTypeInfo() throws Exception { null, null, null)); - MatcherAssert.assertThat(results, is(matchers)); + assertThat(results).isEqualTo(matchers); } } @@ -1465,7 +1456,7 @@ public void testGetTypeInfoWithFiltering() throws Exception { null, "10", null)); - MatcherAssert.assertThat(results, is(matchers)); + assertThat(results).isEqualTo(matchers); } } @@ -1479,28 +1470,30 @@ public void testGetCommandCrossReference() throws Exception { final List> results = getResults(stream); - final List> matchers = + final List> matchers = asList( - nullValue(String.class), // pk_catalog_name - is("APP"), // pk_schema_name - is("FOREIGNTABLE"), // pk_table_name - is("ID"), // pk_column_name - nullValue(String.class), // fk_catalog_name - is("APP"), // fk_schema_name - is("INTTABLE"), // fk_table_name - is("FOREIGNID"), // fk_column_name - is("1"), // key_sequence - containsString("SQL"), // fk_key_name - containsString("SQL"), // pk_key_name - is("3"), // update_rule - is("3")); // delete_rule + new Condition<>(Objects::isNull, "pk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "pk_schema_name expected to equal APP"), + new Condition<>( + c -> c.equals("FOREIGNTABLE"), "pk_table_name should equal FOREIGNTABLE"), + new Condition<>(c -> c.equals("ID"), "pk_column_name should equal ID"), + new Condition<>(Objects::isNull, "fk_catalog_name expected to be null"), + new Condition<>(c -> c.equals("APP"), "fk_schema_name expected to be APP"), + new Condition<>(c -> c.equals("INTTABLE"), "fk_table_name expeced to be INTTABLE"), + new Condition<>( + c -> c.equals("FOREIGNID"), "fk_column_name expected to equal FOREIGNID"), + new Condition<>(c -> c.equals("1"), "key_sequence expected to equal 1"), + new Condition<>(c -> c.contains("SQL"), "fk_key_name expected to contain SQL"), + new Condition<>(c -> c.contains("SQL"), "pk_key_name expected to contain SQL"), + new Condition<>(c -> c.equals("3"), "update_rule expected to equal 3"), + new Condition<>(c -> c.equals("3"), "delete_rule expected to equal 3")); assertEquals(1, results.size()); final List assertions = new ArrayList<>(); for (int i = 0; i < matchers.size(); i++) { final String actual = results.get(0).get(i); - final Matcher expected = matchers.get(i); - assertions.add(() -> MatcherAssert.assertThat(actual, expected)); + final Condition expected = matchers.get(i); + assertions.add(() -> assertThat(actual).satisfies(expected)); } assertAll(assertions); } @@ -1509,7 +1502,7 @@ public void testGetCommandCrossReference() throws Exception { @Test public void testCreateStatementSchema() throws Exception { final FlightInfo info = sqlClient.execute("SELECT * FROM intTable"); - MatcherAssert.assertThat(info.getSchemaOptional(), is(Optional.of(SCHEMA_INT_TABLE))); + assertThat(info.getSchemaOptional()).isEqualTo(Optional.of(SCHEMA_INT_TABLE)); // Consume statement to close connection before cache eviction try (FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { @@ -1526,11 +1519,10 @@ public void testCreateStatementResults() throws Exception { sqlClient.execute("SELECT * FROM intTable").getEndpoints().get(0).getTicket())) { assertAll( () -> { - MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)); + assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE); }, () -> { - MatcherAssert.assertThat( - getResults(stream), is(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY)); + assertThat(getResults(stream)).isEqualTo(EXPECTED_RESULTS_FOR_STAR_SELECT_QUERY); }); } } @@ -1543,19 +1535,19 @@ public void testExecuteUpdate() { sqlClient.executeUpdate( "INSERT INTO INTTABLE (keyName, value) VALUES " + "('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)"); - MatcherAssert.assertThat(insertedCount, is(3L)); + assertThat(insertedCount).isEqualTo(3L); }, () -> { long updatedCount = sqlClient.executeUpdate( "UPDATE INTTABLE SET keyName = 'KEYNAME1' " + "WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'"); - MatcherAssert.assertThat(updatedCount, is(2L)); + assertThat(updatedCount).isEqualTo(2L); }, () -> { long deletedCount = sqlClient.executeUpdate("DELETE FROM INTTABLE WHERE keyName = 'KEYNAME1'"); - MatcherAssert.assertThat(deletedCount, is(3L)); + assertThat(deletedCount).isEqualTo(3L); }); } @@ -1566,10 +1558,10 @@ public void testQueryWithNoResultsShouldNotHang() throws Exception { final FlightStream stream = sqlClient.getStream(preparedStatement.execute().getEndpoints().get(0).getTicket())) { assertAll( - () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)), + () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> { final List> result = getResults(stream); - MatcherAssert.assertThat(result, is(emptyList())); + assertThat(result).isEqualTo(emptyList()); }); } } diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java index 36d621ad64..ee1507b6af 100644 --- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java +++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java @@ -18,7 +18,7 @@ import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults; import static org.apache.arrow.util.AutoCloseables.close; -import static org.hamcrest.CoreMatchers.is; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import org.apache.arrow.flight.FlightClient; @@ -34,7 +34,6 @@ import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; -import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -89,10 +88,10 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce for (FlightEndpoint endpoint : flightInfo.getEndpoints()) { try (FlightStream stream = sqlClient.getStream(endpoint.getTicket())) { assertAll( - () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)), + () -> assertThat(stream.getSchema()).isEqualTo(SCHEMA_INT_TABLE), () -> - MatcherAssert.assertThat( - getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING))); + assertThat(getResults(stream)) + .isEqualTo(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)); } } } diff --git a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java index 71c0dc88e4..3f527f961e 100644 --- a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java +++ b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStreams.java @@ -22,7 +22,7 @@ import static org.apache.arrow.flight.sql.util.FlightStreamUtils.getResults; import static org.apache.arrow.util.AutoCloseables.close; import static org.apache.arrow.vector.types.Types.MinorType.INT; -import static org.hamcrest.CoreMatchers.is; +import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import com.google.common.collect.ImmutableList; @@ -53,7 +53,6 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; -import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -245,15 +244,15 @@ public void testGetTablesResultNoSchema() throws Exception { sqlClient.getTables(null, null, null, null, false).getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA_NO_SCHEMA), () -> { final List> results = getResults(stream); final List> expectedResults = ImmutableList.of( // catalog_name | schema_name | table_name | table_type | table_schema asList(null, null, "test_table", "TABLE")); - MatcherAssert.assertThat(results, is(expectedResults)); + assertThat(results).isEqualTo(expectedResults); }); } } @@ -264,15 +263,15 @@ public void testGetTableTypesResult() throws Exception { sqlClient.getStream(sqlClient.getTableTypes().getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA)), + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA), () -> { final List> tableTypes = getResults(stream); final List> expectedTableTypes = ImmutableList.of( // table_type singletonList("TABLE")); - MatcherAssert.assertThat(tableTypes, is(expectedTableTypes)); + assertThat(tableTypes).isEqualTo(expectedTableTypes); }); } } @@ -283,9 +282,9 @@ public void testGetSqlInfoResults() throws Exception { try (final FlightStream stream = sqlClient.getStream(info.getEndpoints().get(0).getTicket())) { assertAll( () -> - MatcherAssert.assertThat( - stream.getSchema(), is(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA)), - () -> MatcherAssert.assertThat(getResults(stream), is(emptyList()))); + assertThat(stream.getSchema()) + .isEqualTo(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA), + () -> assertThat(getResults(stream)).isEqualTo(emptyList())); } } @@ -303,7 +302,7 @@ public void testGetTypeInfo() throws Exception { "Integer", "4", "400", null, null, "3", "true", null, "true", null, "true", "Integer", null, null, "4", null, "10", null)); - MatcherAssert.assertThat(results, is(matchers)); + assertThat(results).isEqualTo(matchers); } } @@ -317,10 +316,8 @@ public void testExecuteQuery() throws Exception { .get(0) .getTicket())) { assertAll( - () -> - MatcherAssert.assertThat(stream.getSchema(), is(FlightSqlTestProducer.FIXED_SCHEMA)), - () -> - MatcherAssert.assertThat(getResults(stream), is(singletonList(singletonList("1"))))); + () -> assertThat(stream.getSchema()).isEqualTo(FlightSqlTestProducer.FIXED_SCHEMA), + () -> assertThat(getResults(stream)).isEqualTo(singletonList(singletonList("1")))); } } } diff --git a/flight/pom.xml b/flight/pom.xml index 2fc3e89ef8..30f75fa27e 100644 --- a/flight/pom.xml +++ b/flight/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-flight diff --git a/format/pom.xml b/format/pom.xml index d3578b63d2..8c2f2d3387 100644 --- a/format/pom.xml +++ b/format/pom.xml @@ -23,7 +23,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-format diff --git a/gandiva/pom.xml b/gandiva/pom.xml index 5367bfdedf..190bf016ce 100644 --- a/gandiva/pom.xml +++ b/gandiva/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT org.apache.arrow.gandiva diff --git a/gandiva/src/main/cpp/jni_common.cc b/gandiva/src/main/cpp/jni_common.cc index 2851250072..ec4888a512 100644 --- a/gandiva/src/main/cpp/jni_common.cc +++ b/gandiva/src/main/cpp/jni_common.cc @@ -221,7 +221,7 @@ DataTypePtr ProtoTypeToDataType(const gandiva::types::ExtGandivaType& ext_type) return arrow::date64(); case gandiva::types::DECIMAL: // TODO: error handling - return arrow::decimal(ext_type.precision(), ext_type.scale()); + return arrow::decimal128(ext_type.precision(), ext_type.scale()); case gandiva::types::TIME32: return ProtoTypeToTime32(ext_type); case gandiva::types::TIME64: diff --git a/memory/memory-core/pom.xml b/memory/memory-core/pom.xml index 72ee69d60a..825b3dae4b 100644 --- a/memory/memory-core/pom.xml +++ b/memory/memory-core/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-memory - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-memory-core @@ -100,8 +100,8 @@ under the License. test - - + + **/TestOpens.java diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java index 5d052c2cde..d4d76f57f4 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java @@ -16,7 +16,7 @@ */ package org.apache.arrow.memory; -import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicLongFieldUpdater; import org.apache.arrow.util.Preconditions; import org.checkerframework.checker.nullness.qual.Nullable; @@ -37,16 +37,24 @@ class Accountant implements AutoCloseable { */ protected final long reservation; - private final AtomicLong peakAllocation = new AtomicLong(); + // AtomicLongFieldUpdaters for memory accounting fields to reduce memory overhead + private static final AtomicLongFieldUpdater PEAK_ALLOCATION_UPDATER = + AtomicLongFieldUpdater.newUpdater(Accountant.class, "peakAllocation"); + private static final AtomicLongFieldUpdater ALLOCATION_LIMIT_UPDATER = + AtomicLongFieldUpdater.newUpdater(Accountant.class, "allocationLimit"); + private static final AtomicLongFieldUpdater LOCALLY_HELD_MEMORY_UPDATER = + AtomicLongFieldUpdater.newUpdater(Accountant.class, "locallyHeldMemory"); + + private volatile long peakAllocation = 0; /** * Maximum local memory that can be held. This can be externally updated. Changing it won't cause * past memory to change but will change responses to future allocation efforts */ - private final AtomicLong allocationLimit = new AtomicLong(); + private volatile long allocationLimit = 0; /** Currently allocated amount of memory. */ - private final AtomicLong locallyHeldMemory = new AtomicLong(); + private volatile long locallyHeldMemory = 0; public Accountant( @Nullable Accountant parent, String name, long reservation, long maxAllocation) { @@ -64,7 +72,7 @@ public Accountant( this.parent = parent; this.name = name; this.reservation = reservation; - this.allocationLimit.set(maxAllocation); + ALLOCATION_LIMIT_UPDATER.set(this, maxAllocation); if (reservation != 0) { Preconditions.checkArgument(parent != null, "parent must not be null"); @@ -117,12 +125,12 @@ private AllocationOutcome.Status allocateBytesInternal(long size) { } private void updatePeak() { - final long currentMemory = locallyHeldMemory.get(); + final long currentMemory = locallyHeldMemory; while (true) { - final long previousPeak = peakAllocation.get(); + final long previousPeak = peakAllocation; if (currentMemory > previousPeak) { - if (!peakAllocation.compareAndSet(previousPeak, currentMemory)) { + if (!PEAK_ALLOCATION_UPDATER.compareAndSet(this, previousPeak, currentMemory)) { // peak allocation changed underneath us. try again. continue; } @@ -166,7 +174,7 @@ private AllocationOutcome.Status allocate( final boolean incomingUpdatePeak, final boolean forceAllocation, @Nullable AllocationOutcomeDetails details) { - final long oldLocal = locallyHeldMemory.getAndAdd(size); + final long oldLocal = LOCALLY_HELD_MEMORY_UPDATER.getAndAdd(this, size); final long newLocal = oldLocal + size; // Borrowed from Math.addExact (but avoid exception here) // Overflow if result has opposite sign of both arguments @@ -174,7 +182,7 @@ private AllocationOutcome.Status allocate( // failure final boolean overflow = ((oldLocal ^ newLocal) & (size ^ newLocal)) < 0; final long beyondReservation = newLocal - reservation; - final boolean beyondLimit = overflow || newLocal > allocationLimit.get(); + final boolean beyondLimit = overflow || newLocal > allocationLimit; final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit); if (details != null) { @@ -214,7 +222,7 @@ private AllocationOutcome.Status allocate( public void releaseBytes(long size) { // reduce local memory. all memory released above reservation should be released up the tree. - final long newSize = locallyHeldMemory.addAndGet(-size); + final long newSize = LOCALLY_HELD_MEMORY_UPDATER.addAndGet(this, -size); Preconditions.checkArgument(newSize >= 0, "Accounted size went negative."); @@ -255,7 +263,7 @@ public String getName() { * @return Limit in bytes. */ public long getLimit() { - return allocationLimit.get(); + return allocationLimit; } /** @@ -274,7 +282,7 @@ public long getInitReservation() { * @param newLimit The limit in bytes. */ public void setLimit(long newLimit) { - allocationLimit.set(newLimit); + ALLOCATION_LIMIT_UPDATER.set(this, newLimit); } /** @@ -284,7 +292,7 @@ public void setLimit(long newLimit) { * @return Currently allocate memory in bytes. */ public long getAllocatedMemory() { - return locallyHeldMemory.get(); + return locallyHeldMemory; } /** @@ -293,17 +301,17 @@ public long getAllocatedMemory() { * @return The peak allocated memory in bytes. */ public long getPeakMemoryAllocation() { - return peakAllocation.get(); + return peakAllocation; } public long getHeadroom() { - long localHeadroom = allocationLimit.get() - locallyHeldMemory.get(); + long localHeadroom = allocationLimit - locallyHeldMemory; if (parent == null) { return localHeadroom; } // Amount of reserved memory left on top of what parent has - long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory.get()); + long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory); return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom); } } diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java index 775a8925ad..9712be34d7 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java @@ -24,7 +24,6 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.ReadOnlyBufferException; -import java.util.concurrent.atomic.AtomicLong; import org.apache.arrow.memory.BaseAllocator.Verbosity; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.HistoricalLog; @@ -57,9 +56,8 @@ public final class ArrowBuf implements AutoCloseable { private static final int DOUBLE_SIZE = Double.BYTES; private static final int LONG_SIZE = Long.BYTES; - private static final AtomicLong idGenerator = new AtomicLong(0); private static final int LOG_BYTES_PER_ROW = 10; - private final long id = idGenerator.incrementAndGet(); + private final ReferenceManager referenceManager; private final @Nullable BufferManager bufferManager; private final long addr; @@ -67,7 +65,8 @@ public final class ArrowBuf implements AutoCloseable { private long writerIndex; private final @Nullable HistoricalLog historicalLog = BaseAllocator.DEBUG - ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", id) + ? new HistoricalLog( + BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", System.identityHashCode(this)) : null; private volatile long capacity; @@ -136,7 +135,7 @@ public long capacity() { /** * Adjusts the capacity of this buffer. Size increases are NOT supported. * - * @param newCapacity Must be in in the range [0, length). + * @param newCapacity Must be in the range [0, length). */ public synchronized ArrowBuf capacity(long newCapacity) { @@ -218,7 +217,8 @@ public long memoryAddress() { @Override public String toString() { - return String.format("ArrowBuf[%d], address:%d, capacity:%d", id, memoryAddress(), capacity); + return String.format( + "ArrowBuf[%d], address:%d, capacity:%d", getId(), memoryAddress(), capacity); } @Override @@ -1080,12 +1080,15 @@ public String toHexString(final long start, final int length) { } /** - * Get the integer id assigned to this ArrowBuf for debugging purposes. + * Get the id assigned to this ArrowBuf for debugging purposes. + * + *

Returns {@link System#identityHashCode(Object)} which provides a unique identifier for this + * buffer without any per-instance memory overhead. * - * @return integer id + * @return the identity hash code for this buffer */ public long getId() { - return id; + return System.identityHashCode(this); } /** diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java index 4f9d3c61c6..dbd6da3291 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java @@ -25,9 +25,10 @@ public interface BufferAllocator extends AutoCloseable { /** - * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be - * larger than the requested size for rounding purposes. However, the buffer's capacity will be - * set to the configured size. + * Allocate a new or reused buffer of the provided size. The buffer may be larger than the + * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will + * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to + * the requested size if needed. * * @param size The size in bytes. * @return a new ArrowBuf, or null if the request can't be satisfied @@ -36,9 +37,10 @@ public interface BufferAllocator extends AutoCloseable { ArrowBuf buffer(long size); /** - * Allocate a new or reused buffer of the provided size. Note that the buffer may technically be - * larger than the requested size for rounding purposes. However, the buffer's capacity will be - * set to the configured size. + * Allocate a new or reused buffer of the provided size. The buffer may be larger than the + * requested size for rounding purposes (e.g. to a power of two), and the buffer's capacity will + * reflect the actual allocated size. Use {@link ArrowBuf#capacity(long)} to set the capacity to + * the requested size if needed. * * @param size The size in bytes. * @param manager A buffer manager to manage reallocation. diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java index b562a421e7..eb90efcbb5 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java @@ -17,8 +17,7 @@ package org.apache.arrow.memory; import java.util.IdentityHashMap; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.HistoricalLog; import org.apache.arrow.util.Preconditions; @@ -32,12 +31,13 @@ public class BufferLedger implements ValueWithKeyIncluded, ReferenceManager { private final @Nullable IdentityHashMap buffers = BaseAllocator.DEBUG ? new IdentityHashMap<>() : null; - private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0); - // unique ID assigned to each ledger - private final long ledgerId = LEDGER_ID_GENERATOR.incrementAndGet(); - private final AtomicInteger bufRefCnt = new AtomicInteger(0); // start at zero so we can - // manage request for retain - // correctly + + // AtomicIntegerFieldUpdater for bufRefCnt to reduce memory overhead + private static final AtomicIntegerFieldUpdater BUF_REF_CNT_UPDATER = + AtomicIntegerFieldUpdater.newUpdater(BufferLedger.class, "bufRefCnt"); + // start at zero so we can manage request for retain correctly + private volatile int bufRefCnt = 0; + private final long lCreationTime = System.nanoTime(); private final BufferAllocator allocator; private final AllocationManager allocationManager; @@ -78,7 +78,7 @@ public BufferAllocator getAllocator() { */ @Override public int getRefCount() { - return bufRefCnt.get(); + return bufRefCnt; } /** @@ -86,7 +86,7 @@ public int getRefCount() { * ArrowBufs managed by this ledger will share the ref count. */ void increment() { - bufRefCnt.incrementAndGet(); + BUF_REF_CNT_UPDATER.incrementAndGet(this); } /** @@ -144,7 +144,7 @@ private int decrement(int decrement) { allocator.assertOpen(); final int outcome; synchronized (allocationManager) { - outcome = bufRefCnt.addAndGet(-decrement); + outcome = BUF_REF_CNT_UPDATER.addAndGet(this, -decrement); if (outcome == 0) { lDestructionTime = System.nanoTime(); // refcount of this reference manager has dropped to 0 @@ -174,7 +174,7 @@ public void retain(int increment) { if (historicalLog != null) { historicalLog.recordEvent("retain(%d)", increment); } - final int originalReferenceCount = bufRefCnt.getAndAdd(increment); + final int originalReferenceCount = BUF_REF_CNT_UPDATER.getAndAdd(this, increment); Preconditions.checkArgument(originalReferenceCount > 0); } @@ -472,13 +472,13 @@ public long getAccountedSize() { void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) { CommonUtil.indent(sb, indent) .append("ledger[") - .append(ledgerId) + .append(System.identityHashCode(this)) .append("] allocator: ") .append(allocator.getName()) .append("), isOwning: ") .append(", size: ") .append(", references: ") - .append(bufRefCnt.get()) + .append(bufRefCnt) .append(", life: ") .append(lCreationTime) .append("..") diff --git a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java index 91bd7cd905..be0749a215 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java @@ -18,6 +18,7 @@ import java.lang.reflect.Constructor; import java.lang.reflect.Field; +import java.lang.reflect.InaccessibleObjectException; import java.lang.reflect.InvocationTargetException; import java.nio.ByteBuffer; import java.nio.ByteOrder; @@ -81,9 +82,18 @@ public Object run() { BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class); // get the offset of the address field in a java.nio.Buffer object + long maybeOffset; Field addressField = java.nio.Buffer.class.getDeclaredField("address"); - addressField.setAccessible(true); - BYTE_BUFFER_ADDRESS_OFFSET = UNSAFE.objectFieldOffset(addressField); + try { + addressField.setAccessible(true); + maybeOffset = UNSAFE.objectFieldOffset(addressField); + } catch (InaccessibleObjectException e) { + maybeOffset = -1; + logger.debug( + "Cannot access the address field of java.nio.Buffer. DirectBuffer operations wont be available", + e); + } + BYTE_BUFFER_ADDRESS_OFFSET = maybeOffset; Constructor directBufferConstructor; long address = -1; @@ -109,6 +119,9 @@ public Object run() { } catch (SecurityException e) { logger.debug("Cannot get constructor for direct buffer allocation", e); return e; + } catch (InaccessibleObjectException e) { + logger.debug("Cannot get constructor for direct buffer allocation", e); + return e; } } }); @@ -156,7 +169,11 @@ public Object run() { * @return address of the underlying memory. */ public static long getByteBufferAddress(ByteBuffer buf) { - return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET); + if (BYTE_BUFFER_ADDRESS_OFFSET != -1) { + return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET); + } + throw new UnsupportedOperationException( + "Byte buffer address cannot be obtained because sun.misc.Unsafe or java.nio.DirectByteBuffer.(long, int) is not available"); } private MemoryUtil() {} diff --git a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java index a39004a9d0..ba5a539a87 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java @@ -22,7 +22,9 @@ import java.util.Collections; import java.util.Iterator; import java.util.List; +import java.util.stream.Stream; import java.util.stream.StreamSupport; +import org.checkerframework.checker.nullness.qual.Nullable; /** Utilities for AutoCloseable classes. */ public final class AutoCloseables { @@ -33,7 +35,8 @@ private AutoCloseables() {} * Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on autoCloseables * when close is called. */ - public static AutoCloseable all(final Collection autoCloseables) { + public static AutoCloseable all( + final @Nullable Collection autoCloseables) { return new AutoCloseable() { @Override public void close() throws Exception { @@ -48,7 +51,10 @@ public void close() throws Exception { * @param t the throwable to add suppressed exception to * @param autoCloseables the closeables to close */ - public static void close(Throwable t, AutoCloseable... autoCloseables) { + public static void close(Throwable t, @Nullable AutoCloseable... autoCloseables) { + if (autoCloseables == null) { + return; + } close(t, Arrays.asList(autoCloseables)); } @@ -58,7 +64,8 @@ public static void close(Throwable t, AutoCloseable... autoCloseables) { * @param t the throwable to add suppressed exception to * @param autoCloseables the closeables to close */ - public static void close(Throwable t, Iterable autoCloseables) { + public static void close( + Throwable t, @Nullable Iterable autoCloseables) { try { close(autoCloseables); } catch (Exception e) { @@ -71,7 +78,10 @@ public static void close(Throwable t, Iterable autoClos * * @param autoCloseables the closeables to close */ - public static void close(AutoCloseable... autoCloseables) throws Exception { + public static void close(@Nullable AutoCloseable... autoCloseables) throws Exception { + if (autoCloseables == null) { + return; + } close(Arrays.asList(autoCloseables)); } @@ -80,7 +90,8 @@ public static void close(AutoCloseable... autoCloseables) throws Exception { * * @param ac the closeables to close */ - public static void close(Iterable ac) throws Exception { + public static void close(@Nullable Iterable ac) + throws Exception { // this method can be called on a single object if it implements Iterable // like for example VectorContainer make sure we handle that properly if (ac == null) { @@ -111,12 +122,17 @@ public static void close(Iterable ac) throws Exception /** Calls {@link #close(Iterable)} on the flattened list of closeables. */ @SafeVarargs - public static void close(Iterable... closeables) throws Exception { + public static void close(@Nullable Iterable... closeables) + throws Exception { + if (closeables == null) { + return; + } close(flatten(closeables)); } @SafeVarargs - private static Iterable flatten(Iterable... closeables) { + private static Iterable flatten( + Iterable... closeables) { return new Iterable() { // Cast from Iterable to Iterable is safe in this // context @@ -127,16 +143,18 @@ public Iterator iterator() { return Arrays.stream(closeables) .flatMap( (Iterable i) -> - StreamSupport.stream( - ((Iterable) i).spliterator(), /* parallel= */ false)) + i == null + ? Stream.empty() + : StreamSupport.stream( + ((Iterable) i).spliterator(), /* parallel= */ false)) .iterator(); } }; } /** Converts ac to a {@link Iterable} filtering out any null values. */ - public static Iterable iter(AutoCloseable... ac) { - if (ac.length == 0) { + public static Iterable iter(@Nullable AutoCloseable... ac) { + if (ac == null || ac.length == 0) { return Collections.emptyList(); } else { final List nonNullAc = new ArrayList<>(); @@ -153,10 +171,11 @@ public static Iterable iter(AutoCloseable... ac) { public static class RollbackCloseable implements AutoCloseable { private boolean commit = false; - private List closeables; + private final List closeables; - public RollbackCloseable(AutoCloseable... closeables) { - this.closeables = new ArrayList<>(Arrays.asList(closeables)); + public RollbackCloseable(@Nullable AutoCloseable... closeables) { + this.closeables = + closeables == null ? new ArrayList<>() : new ArrayList<>(Arrays.asList(closeables)); } public T add(T t) { @@ -165,12 +184,18 @@ public T add(T t) { } /** Add all of list to the rollback list. */ - public void addAll(AutoCloseable... list) { + public void addAll(@Nullable AutoCloseable... list) { + if (list == null) { + return; + } closeables.addAll(Arrays.asList(list)); } /** Add all of list to the rollback list. */ - public void addAll(Iterable list) { + public void addAll(@Nullable Iterable list) { + if (list == null) { + return; + } for (AutoCloseable ac : list) { closeables.add(ac); } @@ -189,7 +214,7 @@ public void close() throws Exception { } /** Creates an {@link RollbackCloseable} from the given closeables. */ - public static RollbackCloseable rollbackable(AutoCloseable... closeables) { + public static RollbackCloseable rollbackable(@Nullable AutoCloseable... closeables) { return new RollbackCloseable(closeables); } @@ -203,7 +228,7 @@ public static RollbackCloseable rollbackable(AutoCloseable... closeables) { * @throws RuntimeException if an Exception occurs; the Exception is wrapped by the * RuntimeException */ - public static void closeNoChecked(final AutoCloseable autoCloseable) { + public static void closeNoChecked(final @Nullable AutoCloseable autoCloseable) { if (autoCloseable != null) { try { autoCloseable.close(); diff --git a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java index b5e0a71e7e..f74bf63f82 100644 --- a/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java +++ b/memory/memory-core/src/test/java/org/apache/arrow/memory/TestOpens.java @@ -20,32 +20,27 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.condition.JRE.JAVA_16; +import org.apache.arrow.memory.util.MemoryUtil; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledForJreRange; public class TestOpens { - /** Instantiating the RootAllocator should poke MemoryUtil and fail. */ + /** Accessing MemoryUtil.directBuffer should fail as add-opens is not configured. */ @Test @EnabledForJreRange(min = JAVA_16) public void testMemoryUtilFailsLoudly() { // This test is configured by Maven to run WITHOUT add-opens. So this should fail on JDK16+ // (where JEP396 means that add-opens is required to access JDK internals). // The test will likely fail in your IDE if it doesn't correctly pick this up. - Throwable e = - assertThrows( - Throwable.class, - () -> { - BufferAllocator allocator = new RootAllocator(); - allocator.close(); - }); + Throwable e = assertThrows(Throwable.class, () -> MemoryUtil.directBuffer(0, 10)); boolean found = false; while (e != null) { - e = e.getCause(); - if (e instanceof RuntimeException - && e.getMessage().contains("Failed to initialize MemoryUtil")) { + if (e instanceof UnsupportedOperationException + && e.getMessage().contains("java.nio.DirectByteBuffer.(long, int) not available")) { found = true; break; } + e = e.getCause(); } assertTrue(found, "Expected exception was not thrown"); } diff --git a/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java new file mode 100644 index 0000000000..ba5b78178a --- /dev/null +++ b/memory/memory-core/src/test/java/org/apache/arrow/util/TestAutoCloseables.java @@ -0,0 +1,268 @@ +/* + * 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.arrow.util; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +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 java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import org.junit.jupiter.api.Test; + +public class TestAutoCloseables { + + /** Closeable that records that it was closed and can optionally throw. */ + private static final class TrackCloseable implements AutoCloseable { + private boolean closed; + private final Exception toThrow; + + TrackCloseable() { + this.toThrow = null; + } + + TrackCloseable(Exception toThrow) { + this.toThrow = toThrow; + } + + @Override + public void close() throws Exception { + closed = true; + if (toThrow != null) { + throw toThrow; + } + } + + boolean isClosed() { + return closed; + } + } + + @Test + public void testCloseVarargsIgnoresNulls() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + AutoCloseables.close(a, null, b); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + } + + @Test + public void testCloseVarargsThrowsFirstExceptionAndSuppressesRest() throws Exception { + Exception e1 = new Exception("first"); + Exception e2 = new Exception("second"); + TrackCloseable c1 = new TrackCloseable(e1); + TrackCloseable c2 = new TrackCloseable(e2); + Exception thrown = assertThrows(Exception.class, () -> AutoCloseables.close(c1, c2)); + assertEquals("first", thrown.getMessage()); + assertTrue(Arrays.asList(thrown.getSuppressed()).contains(e2)); + } + + @Test + public void testCloseIterableNullIterableReturns() throws Exception { + AutoCloseables.close((List) null); // no exception + } + + @Test + public void testCloseIterableIgnoresNullElements() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + List list = Arrays.asList(a, null, b); + AutoCloseables.close(list); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + } + + @Test + public void testCloseIterableWhenIterableIsAlsoAutoCloseable() throws Exception { + TrackCloseable iter = new TrackCloseable(); + TrackCloseable inner = new TrackCloseable(); + // When the Iterable itself implements AutoCloseable (e.g. VectorContainer), + // close(Iterable) calls close() on it and does not iterate over elements + class IterableCloseable implements Iterable, AutoCloseable { + @Override + @SuppressWarnings("unchecked") + public Iterator iterator() { + return (Iterator) Collections.singletonList(inner); + } + + @Override + public void close() throws Exception { + iter.close(); + } + } + AutoCloseables.close(new IterableCloseable()); + assertTrue(iter.isClosed()); + assertFalse(inner.isClosed()); + } + + @Test + public void testCloseIterableVarargsWithNullIterables() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + TrackCloseable c = new TrackCloseable(); + List list1 = Arrays.asList(null, a, b); + List list2 = Collections.singletonList(c); + AutoCloseables.close(list1, null, list2); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + assertTrue(c.isClosed()); + } + + @Test + public void testCloseThrowableSuppressesException() { + Exception e = new Exception("from close"); + TrackCloseable c = new TrackCloseable(e); + Exception main = new Exception("main"); + AutoCloseables.close(main, c); + assertTrue(c.isClosed()); + assertEquals(1, main.getSuppressed().length); + assertEquals(e, main.getSuppressed()[0]); + } + + @Test + public void testCloseThrowableWithNullCloseables() { + Exception main = new Exception("main"); + AutoCloseables.close(main, (AutoCloseable) null); + assertEquals(0, main.getSuppressed().length); + + AutoCloseables.close(main, (AutoCloseable[]) null); // no exception + } + + @Test + public void testIterFiltersNulls() { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + Iterable it = AutoCloseables.iter(a, null, b); + List list = new ArrayList<>(); + it.forEach(list::add); + assertEquals(2, list.size()); + assertTrue(list.contains(a)); + assertTrue(list.contains(b)); + } + + @Test + public void testIterEmptyVarargs() { + Iterable it = AutoCloseables.iter(); + List list = new ArrayList<>(); + it.forEach(list::add); + assertTrue(list.isEmpty()); + } + + @Test + public void testIterWithNull() { + AutoCloseables.iter((AutoCloseable) null); // no exception + } + + @Test + public void testCloseNoCheckedWithNull() { + AutoCloseables.closeNoChecked(null); // no exception + } + + @Test + public void testCloseNoCheckedWrapsException() { + Exception e = new Exception("close failed"); + TrackCloseable c = new TrackCloseable(e); + RuntimeException re = + assertThrows(RuntimeException.class, () -> AutoCloseables.closeNoChecked(c)); + assertSame(re.getCause(), e); + assertTrue(re.getMessage().contains("close failed")); + } + + @Test + public void testNoop() throws Exception { + AutoCloseable noop = AutoCloseables.noop(); + assertSame(noop, AutoCloseables.noop()); + noop.close(); // no exception + } + + @Test + public void testAllClosesCollectionOnClose() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + List list = Arrays.asList(a, b); + AutoCloseable all = AutoCloseables.all(list); + assertFalse(a.isClosed()); + assertFalse(b.isClosed()); + all.close(); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + } + + @Test + public void testAllWithNullCollection() throws Exception { + AutoCloseable all = AutoCloseables.all(null); + all.close(); // no exception + } + + @Test + public void testRollbackCloseableClosesWhenNotCommitted() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b); + rb.close(); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + } + + @Test + public void testRollbackCloseableDoesNotCloseWhenCommitted() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, b); + rb.commit(); + rb.close(); + assertFalse(a.isClosed()); + assertFalse(b.isClosed()); + } + + @Test + public void testRollbackCloseableAddAndAddAll() throws Exception { + TrackCloseable a = new TrackCloseable(); + TrackCloseable b = new TrackCloseable(); + TrackCloseable c = new TrackCloseable(); + TrackCloseable d = new TrackCloseable(); + AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a); + rb.add(b); + rb.addAll(c, d); + rb.addAll((AutoCloseable[]) null); // null varargs shouldn't fail + rb.addAll((List) null); // null Iterable shouldn't fail + rb.close(); + assertTrue(a.isClosed()); + assertTrue(b.isClosed()); + assertTrue(c.isClosed()); + assertTrue(d.isClosed()); + } + + @Test + public void testRollbackCloseableWithNull() throws Exception { + AutoCloseables.rollbackable((AutoCloseable) null); // no exception + } + + @Test + public void testRollbackCloseableWithNulls() throws Exception { + TrackCloseable a = new TrackCloseable(); + AutoCloseables.RollbackCloseable rb = AutoCloseables.rollbackable(a, null); + rb.close(); + assertTrue(a.isClosed()); + } +} diff --git a/memory/memory-netty-buffer-patch/pom.xml b/memory/memory-netty-buffer-patch/pom.xml index 07dc7d2403..039b2aa04a 100644 --- a/memory/memory-netty-buffer-patch/pom.xml +++ b/memory/memory-netty-buffer-patch/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-memory - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-memory-netty-buffer-patch diff --git a/memory/memory-netty/pom.xml b/memory/memory-netty/pom.xml index 6d660da117..4218910980 100644 --- a/memory/memory-netty/pom.xml +++ b/memory/memory-netty/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-memory - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-memory-netty diff --git a/memory/memory-unsafe/pom.xml b/memory/memory-unsafe/pom.xml index 92dc0c9fe5..3fafb42802 100644 --- a/memory/memory-unsafe/pom.xml +++ b/memory/memory-unsafe/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-memory - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-memory-unsafe diff --git a/memory/pom.xml b/memory/pom.xml index bc34c26050..4ea3d1f9ca 100644 --- a/memory/pom.xml +++ b/memory/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-memory pom diff --git a/performance/pom.xml b/performance/pom.xml index 3f18188e3a..d994bff45b 100644 --- a/performance/pom.xml +++ b/performance/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-performance jar @@ -35,10 +35,10 @@ under the License. true .* 1 - + 5 5 - + jmh-result.json json @@ -75,7 +75,7 @@ under the License. com.h2database h2 - 2.3.232 + 2.4.240 runtime @@ -143,7 +143,7 @@ under the License. java -classpath - + org.openjdk.jmh.Main ${benchmark.filter} -f diff --git a/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java new file mode 100644 index 0000000000..395ba13b9d --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/memory/MemoryFootprintBenchmarks.java @@ -0,0 +1,213 @@ +/* + * 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.arrow.memory; + +import java.lang.management.ManagementFactory; +import java.lang.management.MemoryMXBean; +import java.lang.management.MemoryUsage; +import java.util.concurrent.TimeUnit; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.RunnerException; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; + +/** + * Benchmarks for memory footprint of Arrow memory objects. + * + *

This benchmark measures the heap memory overhead of creating many ArrowBuf instances. The + * optimizations using AtomicFieldUpdater instead of AtomicLong/AtomicInteger objects should reduce + * memory overhead significantly. + * + *

Expected savings per instance: - ArrowBuf: 8 bytes (id field removed) - BufferLedger: 28 bytes + * (20 from AtomicInteger + 8 from ledgerId) - Accountant: 48 bytes (3 × 16 bytes from AtomicLong + * objects) + * + *

For 1M ArrowBuf instances, this should save approximately 8 MB of heap memory. + */ +@State(Scope.Benchmark) +@Fork( + value = 1, + jvmArgs = {"-Xms2g", "-Xmx2g"}) +@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS) +public class MemoryFootprintBenchmarks { + + /** Number of ArrowBuf instances to create for memory footprint measurement. */ + private static final int NUM_BUFFERS = 100_000; + + /** Size in bytes of each buffer allocation. */ + private static final int BUFFER_SIZE = 1024; + + /** Root allocator used for all buffer allocations in the benchmark. */ + private RootAllocator allocator; + + /** Array to hold references to allocated buffers, preventing garbage collection. */ + private ArrowBuf[] buffers; + + /** JMX bean for querying heap memory usage statistics. */ + private MemoryMXBean memoryBean; + + /** + * Sets up the benchmark state before each trial. + * + *

Initializes the memory monitoring bean, creates a root allocator with sufficient capacity, + * and allocates the buffer reference array. + */ + @Setup(Level.Trial) + public void setup() { + memoryBean = ManagementFactory.getMemoryMXBean(); + allocator = new RootAllocator((long) NUM_BUFFERS * BUFFER_SIZE); + buffers = new ArrowBuf[NUM_BUFFERS]; + } + + /** + * Cleans up buffers after each benchmark invocation. + * + *

Closes all allocated buffers to prevent memory leaks and ensure each iteration starts with a + * clean slate. This is critical for the memory footprint benchmark which allocates many buffers + * that would otherwise accumulate across warmup and measurement iterations. + */ + @TearDown(Level.Invocation) + public void tearDown() { + for (int i = 0; i < NUM_BUFFERS; i++) { + if (buffers[i] != null) { + buffers[i].close(); + buffers[i] = null; + } + } + } + + /** + * Cleans up the allocator after the trial completes. + * + *

Closes the root allocator to release all resources after all warmup and measurement + * iterations are complete. + */ + @TearDown(Level.Trial) + public void tearDownTrial() { + allocator.close(); + } + + /** + * Benchmark that measures heap memory usage when creating many ArrowBuf instances. + * + *

This benchmark creates {@value #NUM_BUFFERS} ArrowBuf instances and measures the heap memory + * used. With the AtomicFieldUpdater optimizations, we expect to save approximately 800 KB of heap + * memory (8 bytes × 100,000 instances) just from removing the id field in ArrowBuf. + * + *

The benchmark performs garbage collection before and after allocation to ensure accurate + * measurement of heap memory delta. Results are printed to stdout for analysis. + * + * @return the total heap memory used by the allocated buffers in bytes + */ + @Benchmark + @BenchmarkMode(Mode.SingleShotTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public long measureArrowBufMemoryFootprint() { + // Force GC before measurement + System.gc(); + System.gc(); + System.gc(); + + MemoryUsage heapBefore = memoryBean.getHeapMemoryUsage(); + long usedBefore = heapBefore.getUsed(); + + // Allocate buffers + for (int i = 0; i < NUM_BUFFERS; i++) { + buffers[i] = allocator.buffer(BUFFER_SIZE); + } + + // Force GC to get accurate measurement + System.gc(); + System.gc(); + System.gc(); + + MemoryUsage heapAfter = memoryBean.getHeapMemoryUsage(); + long usedAfter = heapAfter.getUsed(); + + long memoryUsed = usedAfter - usedBefore; + + // Print memory usage for analysis + System.out.printf( + "Created %d ArrowBuf instances. Heap memory used: %d bytes (%.2f MB)%n", + NUM_BUFFERS, memoryUsed, memoryUsed / (1024.0 * 1024.0)); + System.out.printf( + "Average memory per ArrowBuf: %.2f bytes%n", (double) memoryUsed / NUM_BUFFERS); + + return memoryUsed; + } + + /** + * Benchmark that measures allocation and deallocation performance. + * + *

This complements the memory footprint benchmark by measuring the time it takes to allocate + * and deallocate 1,000 buffers in a tight loop. This helps identify any performance regressions + * introduced by memory optimizations. + * + *

Uses a local buffer array to avoid interference with the shared {@link #buffers} array used + * by other benchmarks. + */ + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void measureAllocationPerformance() { + ArrowBuf[] localBuffers = new ArrowBuf[1000]; + + for (int i = 0; i < 1000; i++) { + localBuffers[i] = allocator.buffer(BUFFER_SIZE); + } + + for (int i = 0; i < 1000; i++) { + localBuffers[i].close(); + } + } + + /** + * Main entry point for running the benchmarks standalone. + * + *

This allows running the benchmarks directly from the command line or IDE without using the + * Maven JMH plugin. Example usage: + * + * 

{@code
+   * java -cp target/benchmarks.jar org.apache.arrow.memory.MemoryFootprintBenchmarks
+   * }
+ * + * @param args command line arguments (not used) + * @throws RunnerException if the benchmark runner encounters an error + */ + public static void main(String[] args) throws RunnerException { + Options opt = + new OptionsBuilder() + .include(MemoryFootprintBenchmarks.class.getSimpleName()) + .forks(1) + .build(); + + new Runner(opt).run(); + } +} diff --git a/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java new file mode 100644 index 0000000000..b5f87e7a75 --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/vector/UuidVectorBenchmarks.java @@ -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.arrow.vector; + +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.complex.impl.UuidWriterImpl; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.profile.GCProfiler; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.RunnerException; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; + +/** Benchmarks for {@link UuidVector}. */ +@State(Scope.Benchmark) +public class UuidVectorBenchmarks { + // checkstyle:off: MissingJavadocMethod + + private static final int VECTOR_LENGTH = 10_000; + + private static final int ALLOCATOR_CAPACITY = 1024 * 1024; + + private BufferAllocator allocator; + + private UuidVector vector; + + private UUID[] testUuids; + + @Setup + public void prepare() { + allocator = new RootAllocator(ALLOCATOR_CAPACITY); + vector = new UuidVector("vector", allocator); + vector.allocateNew(VECTOR_LENGTH); + vector.setValueCount(VECTOR_LENGTH); + + // Pre-generate UUIDs for consistent benchmarking + testUuids = new UUID[VECTOR_LENGTH]; + for (int i = 0; i < VECTOR_LENGTH; i++) { + testUuids[i] = new UUID(i, i * 2L); + } + } + + @TearDown + public void tearDown() { + vector.close(); + allocator.close(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void setWithHolder() { + NullableUuidHolder holder = new NullableUuidHolder(); + for (int i = 0; i < VECTOR_LENGTH; i++) { + vector.get(i, holder); + vector.setSafe(i, holder); + } + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void setUuidDirectly() { + for (int i = 0; i < VECTOR_LENGTH; i++) { + vector.setSafe(i, testUuids[i]); + } + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void setWithWriter() { + UuidWriterImpl writer = new UuidWriterImpl(vector); + for (int i = 0; i < VECTOR_LENGTH; i++) { + writer.writeExtension(testUuids[i]); + } + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void getWithUuidHolder() { + NullableUuidHolder holder = new NullableUuidHolder(); + for (int i = 0; i < VECTOR_LENGTH; i++) { + vector.get(i, holder); + } + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MICROSECONDS) + public void getUuidDirectly() { + for (int i = 0; i < VECTOR_LENGTH; i++) { + UUID uuid = vector.getObject(i); + } + } + + public static void main(String[] args) throws RunnerException { + Options opt = + new OptionsBuilder() + .include(UuidVectorBenchmarks.class.getSimpleName()) + .forks(1) + .addProfiler(GCProfiler.class) + .build(); + + new Runner(opt).run(); + } + // checkstyle:on: MissingJavadocMethod +} diff --git a/pom.xml b/pom.xml index ab988c7d05..78137eb4e9 100644 --- a/pom.xml +++ b/pom.xml @@ -23,12 +23,12 @@ under the License. org.apache apache - 34 + 38 org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT pom Apache Arrow Java Root POM @@ -68,6 +68,7 @@ under the License. bom format memory + arrow-variant vector tools adapter/jdbc @@ -91,38 +92,42 @@ under the License. + 1773644827 ${project.build.directory}/generated-sources 1.9.0 - 5.12.1 - 2.0.17 - 33.4.5-jre - 4.1.119.Final - 1.71.0 - 4.30.1 - 2.18.3 - 3.4.1 + 6.1.1 + 2.0.18 + 33.6.0-jre + 4.2.15.Final + 1.82.1 + 4.35.1 + 2.22.0 + 3.5.0 25.2.10 - 1.12.0 - 5.17.0 - + 1.12.1 + 1.17.1 + 5.23.0 + 2 10.23.0 true - 2.37.0 - 3.49.2 - 1.5.18 + 2.42.0 + 4.2.1 + 1.5.37 none -Xdoclint:none --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED - 11 - 11 - 11 - 11 + 17 + 17 + 17 + 17 3.2.2 @@ -172,13 +177,13 @@ under the License. org.assertj assertj-core - 3.27.3 + 3.27.7 test org.immutables value-annotations - 2.10.1 + 2.12.2 provided @@ -309,7 +314,7 @@ under the License. org.immutables value - 2.10.1 + 2.12.2 @@ -347,7 +352,7 @@ under the License. org.jacoco jacoco-maven-plugin - 0.8.12 + 0.8.15 + codegen + ${basedir}/src/main/codegen + + **/*.tdd + **/*.fmpp + **/*.ftl + + + + + diff --git a/vector/src/main/codegen/includes/vv_imports.ftl b/vector/src/main/codegen/includes/vv_imports.ftl index 7f216a7b43..2bbcecc856 100644 --- a/vector/src/main/codegen/includes/vv_imports.ftl +++ b/vector/src/main/codegen/includes/vv_imports.ftl @@ -34,6 +34,7 @@ import org.apache.arrow.vector.complex.*; import org.apache.arrow.vector.complex.reader.*; import org.apache.arrow.vector.complex.impl.*; import org.apache.arrow.vector.complex.writer.*; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter; diff --git a/vector/src/main/codegen/templates/AbstractFieldReader.java b/vector/src/main/codegen/templates/AbstractFieldReader.java index 25b071fab7..789295e959 100644 --- a/vector/src/main/codegen/templates/AbstractFieldReader.java +++ b/vector/src/main/codegen/templates/AbstractFieldReader.java @@ -29,9 +29,9 @@ * Source code generated using FreeMarker template ${.template_name} */ @SuppressWarnings("unused") -abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{ +public abstract class AbstractFieldReader extends AbstractBaseReader implements FieldReader{ - AbstractFieldReader(){ + protected AbstractFieldReader(){ super(); } @@ -108,6 +108,23 @@ public void copyAsField(String name, ${name}Writer writer) { } + + public void read(ExtensionHolder holder) { + fail("Extension"); + } + + public void read(int arrayIndex, ExtensionHolder holder) { + fail("RepeatedExtension"); + } + + public void copyAsValue(AbstractExtensionTypeWriter writer) { + fail("CopyAsValueExtension"); + } + + public void copyAsField(String name, AbstractExtensionTypeWriter writer) { + fail("CopyAsFieldExtension"); + } + public FieldReader reader(String name) { fail("reader(String name)"); return null; @@ -126,4 +143,5 @@ public int size() { private void fail(String name) { throw new IllegalArgumentException(String.format("You tried to read a [%s] type when you are using a field reader of type [%s].", name, this.getClass().getSimpleName())); } + } diff --git a/vector/src/main/codegen/templates/AbstractFieldWriter.java b/vector/src/main/codegen/templates/AbstractFieldWriter.java index ae5b97faef..4b4a17d932 100644 --- a/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -107,14 +107,17 @@ public void endEntry() { throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName())); } + @Override public void write(ExtensionHolder var1) { - this.fail("ExtensionType"); + this.fail("Cannot write ExtensionHolder"); } + @Override public void writeExtension(Object var1) { - this.fail("ExtensionType"); + this.fail("Cannot write extension object"); } - public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) { - this.fail("ExtensionType"); + @Override + public void writeExtension(Object var1, ArrowType type) { + this.fail("Cannot write extension with type " + type); } <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> diff --git a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java index 951edd5eee..2e7792fcfe 100644 --- a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java +++ b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java @@ -295,7 +295,7 @@ public MapWriter map(boolean keysSorted) { @Override public ExtensionWriter extension(ArrowType arrowType) { - return getWriter(MinorType.EXTENSIONTYPE).extension(arrowType); + return getWriter(MinorType.LIST).extension(arrowType); } @Override @@ -325,7 +325,7 @@ public MapWriter map(String name, boolean keysSorted) { @Override public ExtensionWriter extension(String name, ArrowType arrowType) { - return getWriter(MinorType.EXTENSIONTYPE).extension(name, arrowType); + return getWriter(MinorType.STRUCT).extension(name, arrowType); } <#list vv.types as type><#list type.minor as minor> diff --git a/vector/src/main/codegen/templates/ArrowType.java b/vector/src/main/codegen/templates/ArrowType.java index fd35c1cd2b..b428f09155 100644 --- a/vector/src/main/codegen/templates/ArrowType.java +++ b/vector/src/main/codegen/templates/ArrowType.java @@ -27,8 +27,10 @@ import org.apache.arrow.flatbuf.Type; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.types.*; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonIgnore; @@ -331,6 +333,10 @@ public boolean equals(Object obj) { public T accept(ArrowTypeVisitor visitor) { return visitor.visit(this); } + + public FieldWriter getNewFieldWriter(ValueVector vector) { + throw new UnsupportedOperationException("WriterImpl not yet implemented."); + } } private static final int defaultDecimalBitWidth = 128; diff --git a/vector/src/main/codegen/templates/BaseReader.java b/vector/src/main/codegen/templates/BaseReader.java index e75e8a2974..c52345af21 100644 --- a/vector/src/main/codegen/templates/BaseReader.java +++ b/vector/src/main/codegen/templates/BaseReader.java @@ -73,7 +73,7 @@ public interface RepeatedMapReader extends MapReader{ public interface ScalarReader extends <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Reader, - BaseReader {} + ExtensionReader, BaseReader {} interface ComplexReader{ StructReader rootAsStruct(); diff --git a/vector/src/main/codegen/templates/BaseWriter.java b/vector/src/main/codegen/templates/BaseWriter.java index 78da7fddc3..a4c98d7089 100644 --- a/vector/src/main/codegen/templates/BaseWriter.java +++ b/vector/src/main/codegen/templates/BaseWriter.java @@ -125,11 +125,12 @@ public interface ExtensionWriter extends BaseWriter { void writeExtension(Object value); /** - * Adds the given extension type factory. This factory allows configuring writer implementations for specific ExtensionTypeVector. + * Writes the given extension type value. * - * @param factory the extension type factory to add + * @param value the extension type value to write + * @param type of the extension */ - void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory); + void writeExtension(Object value, ArrowType type); } public interface ScalarWriter extends diff --git a/vector/src/main/codegen/templates/ComplexCopier.java b/vector/src/main/codegen/templates/ComplexCopier.java index 4fff7059a7..6655f6c2a7 100644 --- a/vector/src/main/codegen/templates/ComplexCopier.java +++ b/vector/src/main/codegen/templates/ComplexCopier.java @@ -41,11 +41,8 @@ public class ComplexCopier { * @param input field to read from * @param output field to write to */ - public static void copy(FieldReader input, FieldWriter output) { - writeValue(input, output); - } + public static void copy(FieldReader reader, FieldWriter writer) { - private static void writeValue(FieldReader reader, FieldWriter writer) { final MinorType mt = reader.getMinorType(); switch (mt) { @@ -61,7 +58,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) { FieldReader childReader = reader.reader(); FieldWriter childWriter = getListWriterForReader(childReader, writer); if (childReader.isSet()) { - writeValue(childReader, childWriter); + copy(childReader, childWriter); } else { childWriter.writeNull(); } @@ -79,8 +76,8 @@ private static void writeValue(FieldReader reader, FieldWriter writer) { FieldReader structReader = reader.reader(); if (structReader.isSet()) { writer.startEntry(); - writeValue(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key())); - writeValue(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value())); + copy(mapReader.key(), getMapWriterForReader(mapReader.key(), writer.key())); + copy(mapReader.value(), getMapWriterForReader(mapReader.value(), writer.value())); writer.endEntry(); } else { writer.writeNull(); @@ -99,7 +96,7 @@ private static void writeValue(FieldReader reader, FieldWriter writer) { if (childReader.getMinorType() != Types.MinorType.NULL) { FieldWriter childWriter = getStructWriterForReader(childReader, writer, name); if (childReader.isSet()) { - writeValue(childReader, childWriter); + copy(childReader, childWriter); } else { childWriter.writeNull(); } @@ -110,6 +107,16 @@ private static void writeValue(FieldReader reader, FieldWriter writer) { writer.writeNull(); } break; + case EXTENSIONTYPE: + if (reader.isSet()) { + Object value = reader.readObject(); + if (value != null) { + writer.writeExtension(value, reader.getField().getType()); + } + } else { + writer.writeNull(); + } + break; <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> <#assign fields = minor.fields!type.fields /> <#assign uncappedName = name?uncap_first/> @@ -162,6 +169,9 @@ private static FieldWriter getStructWriterForReader(FieldReader reader, StructWr return (FieldWriter) writer.map(name); case LISTVIEW: return (FieldWriter) writer.listView(name); + case EXTENSIONTYPE: + ExtensionWriter extensionWriter = writer.extension(name, reader.getField().getType()); + return (FieldWriter) extensionWriter; default: throw new UnsupportedOperationException(reader.getMinorType().toString()); } @@ -186,6 +196,9 @@ private static FieldWriter getListWriterForReader(FieldReader reader, ListWriter return (FieldWriter) writer.list(); case LISTVIEW: return (FieldWriter) writer.listView(); + case EXTENSIONTYPE: + ExtensionWriter extensionWriter = writer.extension(reader.getField().getType()); + return (FieldWriter) extensionWriter; default: throw new UnsupportedOperationException(reader.getMinorType().toString()); } @@ -211,6 +224,9 @@ private static FieldWriter getMapWriterForReader(FieldReader reader, MapWriter w return (FieldWriter) writer.listView(); case MAP: return (FieldWriter) writer.map(false); + case EXTENSIONTYPE: + ExtensionWriter extensionWriter = writer.extension(reader.getField().getType()); + return (FieldWriter) extensionWriter; default: throw new UnsupportedOperationException(reader.getMinorType().toString()); } diff --git a/vector/src/main/codegen/templates/DenseUnionWriter.java b/vector/src/main/codegen/templates/DenseUnionWriter.java index 8515b759e6..9aeea5b054 100644 --- a/vector/src/main/codegen/templates/DenseUnionWriter.java +++ b/vector/src/main/codegen/templates/DenseUnionWriter.java @@ -55,7 +55,9 @@ public DenseUnionWriter(DenseUnionVector vector, NullableStructWriterFactory nul public void setPosition(int index) { super.setPosition(index); for (BaseWriter writer : writers) { - writer.setPosition(index); + if (writer != null) { + writer.setPosition(index); + } } } diff --git a/vector/src/main/codegen/templates/HolderReaderImpl.java b/vector/src/main/codegen/templates/HolderReaderImpl.java index 1151ea5d39..cdbb65c4f6 100644 --- a/vector/src/main/codegen/templates/HolderReaderImpl.java +++ b/vector/src/main/codegen/templates/HolderReaderImpl.java @@ -126,7 +126,7 @@ public void read(Nullable${name}Holder h) { <#elseif minor.class == "Duration"> return DurationVector.toDuration(holder.value, holder.unit); <#elseif minor.class == "Bit" > - return new Boolean(holder.value != 0); + return Boolean.valueOf(holder.value != 0); <#elseif minor.class == "Decimal"> byte[] bytes = new byte[${type.width}]; holder.buffer.getBytes(holder.start, bytes, 0, ${type.width}); @@ -151,7 +151,7 @@ public void read(Nullable${name}Holder h) { <#elseif minor.class == "TimeStampNano"> return DateUtility.getLocalDateTimeFromEpochNano(holder.value); <#else> - ${friendlyType} value = new ${friendlyType}(this.holder.value); + ${friendlyType} value = ${friendlyType}.valueOf(this.holder.value); return value; } diff --git a/vector/src/main/codegen/templates/NullReader.java b/vector/src/main/codegen/templates/NullReader.java index 1d77248e96..88e6ea98ea 100644 --- a/vector/src/main/codegen/templates/NullReader.java +++ b/vector/src/main/codegen/templates/NullReader.java @@ -86,6 +86,10 @@ public void read(int arrayIndex, Nullable${name}Holder holder){ } + public void read(ExtensionHolder holder) { + holder.isSet = 0; + } + public int size(){ return 0; } diff --git a/vector/src/main/codegen/templates/PromotableWriter.java b/vector/src/main/codegen/templates/PromotableWriter.java index 8d7d57bb9d..11d34f72c9 100644 --- a/vector/src/main/codegen/templates/PromotableWriter.java +++ b/vector/src/main/codegen/templates/PromotableWriter.java @@ -286,7 +286,7 @@ protected void setWriter(ValueVector v) { writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory); break; case EXTENSIONTYPE: - writer = new UnionExtensionWriter((ExtensionTypeVector) vector); + writer = ((ExtensionType) vector.getField().getType()).getNewFieldWriter(vector); break; default: writer = type.getNewFieldWriter(vector); @@ -541,13 +541,13 @@ public void writeLargeVarChar(String value) { } @Override - public void writeExtension(Object value) { - getWriter(MinorType.EXTENSIONTYPE).writeExtension(value); + public void writeExtension(Object value, ArrowType arrowType) { + getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value, arrowType); } @Override - public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) { - getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory); + public void write(ExtensionHolder holder) { + getWriter(MinorType.EXTENSIONTYPE, holder.type()).write(holder); } @Override diff --git a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java index f6e3f63caf..484199ab2a 100644 --- a/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java +++ b/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java @@ -35,6 +35,10 @@ <#include "/@includes/vv_imports.ftl" /> +<#function is_timestamp_tz type> + <#return type?starts_with("TimeStamp") && type?ends_with("TZ")> + + /* * This class is generated using freemarker and the ${.template_name} template. */ @@ -96,55 +100,30 @@ public void close() throws Exception { public void setPosition(int index) { super.setPosition(index); } - <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> - <#assign fields = minor.fields!type.fields /> - <#assign uncappedName = name?uncap_first/> - <#if uncappedName == "int" ><#assign uncappedName = "integer" /> - <#if !minor.typeParams?? > + <#list vv.types as type><#list type.minor as minor> + <#assign lowerName = minor.class?uncap_first /> + <#if lowerName == "int" ><#assign lowerName = "integer" /> + <#assign upperName = minor.class?upper_case /> @Override - public ${name}Writer ${uncappedName}() { + public ${minor.class}Writer ${lowerName}() { return this; } + <#if minor.typeParams?? > @Override - public ${name}Writer ${uncappedName}(String name) { - structName = name; - return writer.${uncappedName}(name); + public ${minor.class}Writer ${lowerName}(String name<#list minor.typeParams as typeParam>, ${typeParam.type} ${typeParam.name}) { + return writer.${lowerName}(name<#list minor.typeParams as typeParam>, ${typeParam.name}); } - - - @Override - public DecimalWriter decimal() { - return this; - } - - @Override - public DecimalWriter decimal(String name, int scale, int precision) { - return writer.decimal(name, scale, precision); - } - - @Override - public DecimalWriter decimal(String name) { - return writer.decimal(name); - } - @Override - public Decimal256Writer decimal256() { - return this; - } - - @Override - public Decimal256Writer decimal256(String name, int scale, int precision) { - return writer.decimal256(name, scale, precision); + public ${minor.class}Writer ${lowerName}(String name) { + structName = name; + return writer.${lowerName}(name); } - @Override - public Decimal256Writer decimal256(String name) { - return writer.decimal256(name); - } + @Override public StructWriter struct() { @@ -215,87 +194,86 @@ public void end() { } @Override - public void write(DecimalHolder holder) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write(holder); - writer.setPosition(writer.idx() + 1); - } - - @Override - public void write(Decimal256Holder holder) { + public void writeNull() { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.write(holder); - writer.setPosition(writer.idx() + 1); + writer.writeNull(); } + <#list vv.types as type> + <#list type.minor as minor> + <#assign name = minor.class?cap_first /> + <#assign fields = minor.fields!type.fields /> + <#assign uncappedName = name?uncap_first/> @Override - public void writeNull() { + public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeNull(); + writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); + writer.setPosition(writer.idx()+1); } - public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) { + <#if is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal(start, buffer, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write(holder); + writer.setPosition(writer.idx()+1); } - public void writeDecimal(BigDecimal value) { + <#elseif minor.class?starts_with("Decimal")> + @Override + public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal(value); - writer.setPosition(writer.idx() + 1); + writer.write${name}(start, buffer, arrowType); + writer.setPosition(writer.idx()+1); } - public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeBigEndianBytesToDecimal(value, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write(holder); + writer.setPosition(writer.idx()+1); } - public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { + @Override + public void write${name}(BigDecimal value) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal256(start, buffer, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write${name}(value); + writer.setPosition(writer.idx()+1); } - public void writeDecimal256(BigDecimal value) { + @Override + public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType){ if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeDecimal256(value); + writer.writeBigEndianBytesTo${name}(value, arrowType); writer.setPosition(writer.idx() + 1); } - - public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { + <#else> + @Override + public void write(${name}Holder holder) { if (writer.idx() >= (idx() + 1) * listSize) { throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); } - writer.writeBigEndianBytesToDecimal256(value, arrowType); - writer.setPosition(writer.idx() + 1); + writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); + writer.setPosition(writer.idx()+1); } + - - <#list vv.types as type> - <#list type.minor as minor> - <#assign name = minor.class?cap_first /> - <#assign fields = minor.fields!type.fields /> - <#assign uncappedName = name?uncap_first/> - <#if minor.class?ends_with("VarBinary")> + <#if minor.class?ends_with("VarBinary")> @Override public void write${minor.class}(byte[] value) { if (writer.idx() >= (idx() + 1) * listSize) { @@ -349,27 +327,8 @@ public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); } - - - <#if !minor.typeParams?? > - @Override - public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, ) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, ); - writer.setPosition(writer.idx() + 1); - } - - public void write(${name}Holder holder) { - if (writer.idx() >= (idx() + 1) * listSize) { - throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); - } - writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, ); - writer.setPosition(writer.idx() + 1); - } + - } diff --git a/vector/src/main/codegen/templates/UnionListWriter.java b/vector/src/main/codegen/templates/UnionListWriter.java index 9424533f29..394348f029 100644 --- a/vector/src/main/codegen/templates/UnionListWriter.java +++ b/vector/src/main/codegen/templates/UnionListWriter.java @@ -53,6 +53,7 @@ public class Union${listName}Writer extends AbstractFieldWriter { private boolean inStruct = false; private boolean listStarted = false; private String structName; + private ArrowType extensionType; <#if listName == "LargeList" || listName == "LargeListView"> private static final long OFFSET_WIDTH = 8; <#else> @@ -122,8 +123,6 @@ public void setPosition(int index) { <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> <#assign upperName = minor.class?upper_case /> - <#assign capName = minor.class?cap_first /> - <#assign vectName = capName /> @Override public ${minor.class}Writer ${lowerName}() { return this; @@ -203,13 +202,13 @@ public MapWriter map(String name, boolean keysSorted) { @Override public ExtensionWriter extension(ArrowType arrowType) { - writer.extension(arrowType); - return writer; + extensionType = arrowType; + return this; } + @Override public ExtensionWriter extension(String name, ArrowType arrowType) { - ExtensionWriter extensionWriter = writer.extension(name, arrowType); - return extensionWriter; + return writer.extension(name, arrowType); } <#if listName == "LargeList"> @@ -336,14 +335,18 @@ public void writeNull() { @Override public void writeExtension(Object value) { - writer.writeExtension(value); + writer.writeExtension(value, extensionType); + writer.setPosition(writer.idx() + 1); } + @Override - public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) { - writer.addExtensionTypeWriterFactory(var1); + public void writeExtension(Object value, ArrowType type) { + writeExtension(value); } + public void write(ExtensionHolder var1) { writer.write(var1); + writer.setPosition(writer.idx() + 1); } <#list vv.types as type> @@ -365,6 +368,7 @@ public void write(${name}Holder holder) { } <#elseif minor.class?starts_with("Decimal")> + @Override public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) { writer.write${name}(start, buffer, arrowType); writer.setPosition(writer.idx()+1); @@ -376,11 +380,13 @@ public void write(${name}Holder holder) { writer.setPosition(writer.idx()+1); } + @Override public void write${name}(BigDecimal value) { writer.write${name}(value); writer.setPosition(writer.idx()+1); } + @Override public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType){ writer.writeBigEndianBytesTo${name}(value, arrowType); writer.setPosition(writer.idx() + 1); @@ -424,6 +430,7 @@ public void write(${name}Holder holder) { writer.setPosition(writer.idx() + 1); } + @Override public void write${minor.class}(String value) { writer.write${minor.class}(value); writer.setPosition(writer.idx() + 1); diff --git a/vector/src/main/codegen/templates/UnionMapWriter.java b/vector/src/main/codegen/templates/UnionMapWriter.java index 8b2f091215..8bbf6ae0a4 100644 --- a/vector/src/main/codegen/templates/UnionMapWriter.java +++ b/vector/src/main/codegen/templates/UnionMapWriter.java @@ -243,4 +243,27 @@ public ExtensionWriter extension(ArrowType type) { return super.extension(type); } } + + public FixedSizeBinaryWriter fixedSizeBinary(int byteWidth) { + switch (mode) { + case KEY: + return entryWriter.fixedSizeBinary(MapVector.KEY_NAME, byteWidth); + case VALUE: + return entryWriter.fixedSizeBinary(MapVector.VALUE_NAME, byteWidth); + default: + return this; + } + } + + @Override + public FixedSizeBinaryWriter fixedSizeBinary() { + switch (mode) { + case KEY: + return entryWriter.fixedSizeBinary(MapVector.KEY_NAME); + case VALUE: + return entryWriter.fixedSizeBinary(MapVector.VALUE_NAME); + default: + return this; + } + } } diff --git a/vector/src/main/codegen/templates/UnionReader.java b/vector/src/main/codegen/templates/UnionReader.java index 96ad3e1b9b..0edae7ade0 100644 --- a/vector/src/main/codegen/templates/UnionReader.java +++ b/vector/src/main/codegen/templates/UnionReader.java @@ -79,6 +79,10 @@ public void read(int index, UnionHolder holder) { } private FieldReader getReaderForIndex(int index) { + return getReaderForIndex(index, null); + } + + private FieldReader getReaderForIndex(int index, ArrowType type) { int typeValue = data.getTypeValue(index); FieldReader reader = (FieldReader) readers[typeValue]; if (reader != null) { @@ -105,11 +109,26 @@ private FieldReader getReaderForIndex(int index) { + case EXTENSIONTYPE: + if(type == null) { + throw new RuntimeException("Cannot get Extension reader without an ArrowType"); + } + return (FieldReader) getExtension(type); default: throw new UnsupportedOperationException("Unsupported type: " + MinorType.values()[typeValue]); } } + private ExtensionReader extensionReader; + + private ExtensionReader getExtension(ArrowType type) { + if (extensionReader == null) { + extensionReader = data.getExtension(type).getReader(); + extensionReader.setPosition(idx()); + } + return extensionReader; + } + private SingleStructReaderImpl structReader; private StructReader getStruct() { @@ -240,4 +259,8 @@ public FieldReader reader() { public boolean next() { return getReaderForIndex(idx()).next(); } + + public void read(ExtensionHolder holder){ + getReaderForIndex(idx(), holder.type()).read(holder); + } } diff --git a/vector/src/main/codegen/templates/UnionVector.java b/vector/src/main/codegen/templates/UnionVector.java index 67efdf60f7..c706591966 100644 --- a/vector/src/main/codegen/templates/UnionVector.java +++ b/vector/src/main/codegen/templates/UnionVector.java @@ -379,6 +379,22 @@ public MapVector getMap(String name, ArrowType arrowType) { return mapVector; } + private ExtensionTypeVector extensionVector; + + public ExtensionTypeVector getExtension(ArrowType arrowType) { + if (extensionVector == null) { + int vectorCount = internalStruct.size(); + extensionVector = addOrGet(null, MinorType.EXTENSIONTYPE, arrowType, ExtensionTypeVector.class); + if (internalStruct.size() > vectorCount) { + extensionVector.allocateNew(); + if (callBack != null) { + callBack.doWork(); + } + } + } + return extensionVector; + } + public int getTypeValue(int index) { return typeBuffer.getByte(index * TYPE_WIDTH); } @@ -725,6 +741,8 @@ public ValueVector getVectorByType(int typeId, ArrowType arrowType) { return getListView(); case MAP: return getMap(name, arrowType); + case EXTENSIONTYPE: + return getExtension(arrowType); default: throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]); } diff --git a/vector/src/main/codegen/templates/UnionWriter.java b/vector/src/main/codegen/templates/UnionWriter.java index 272edab17c..0db699fd8c 100644 --- a/vector/src/main/codegen/templates/UnionWriter.java +++ b/vector/src/main/codegen/templates/UnionWriter.java @@ -28,6 +28,8 @@ package org.apache.arrow.vector.complex.impl; <#include "/@includes/vv_imports.ftl" /> +import java.util.HashMap; + import org.apache.arrow.vector.complex.writer.BaseWriter; import org.apache.arrow.vector.types.Types.MinorType; @@ -213,8 +215,31 @@ public MapWriter asMap(ArrowType arrowType) { return getMapWriter(arrowType); } + private java.util.Map extensionWriters = new HashMap<>(); + private ExtensionWriter getExtensionWriter(ArrowType arrowType) { - throw new UnsupportedOperationException("ExtensionTypes are not supported yet."); + ExtensionWriter w = extensionWriters.get(arrowType); + if (w == null) { + w = ((ExtensionType) arrowType).getNewFieldWriter(data.getExtension(arrowType)); + w.setPosition(idx()); + extensionWriters.put(arrowType, w); + } + return w; + } + + public void writeExtension(Object value, ArrowType type) { + data.setType(idx(), MinorType.EXTENSIONTYPE); + ExtensionWriter w = getExtensionWriter(type); + w.setPosition(idx()); + w.writeExtension(value); + } + + @Override + public void write(ExtensionHolder holder) { + data.setType(idx(), MinorType.EXTENSIONTYPE); + ExtensionWriter w = getExtensionWriter(holder.type()); + w.setPosition(idx()); + w.write(holder); } BaseWriter getWriter(MinorType minorType) { diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java index 4be55396b7..df1ac74f9b 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseFixedWidthVector.java @@ -49,9 +49,7 @@ public abstract class BaseFixedWidthVector extends BaseValueVector protected final Field field; private int allocationMonitor; - protected ArrowBuf validityBuffer; protected ArrowBuf valueBuffer; - protected int valueCount; /** * Constructs a new instance. @@ -72,6 +70,7 @@ public BaseFixedWidthVector(Field field, final BufferAllocator allocator, final refreshValueCapacity(); } + @Override public int getTypeWidth() { return typeWidth; } @@ -87,7 +86,7 @@ public String getName() { /* TODO: * Once the entire hierarchy has been refactored, move common functions - * like getNullCount(), splitAndTransferValidityBuffer to top level + * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be @@ -342,9 +341,9 @@ private void allocateBytes(int valueCount) { * slice the source buffer so we have to explicitly allocate the validityBuffer of the target * vector. This is unlike the databuffer which we can always slice for the target vector. */ - private void allocateValidityBuffer(final int validityBufferSize) { - validityBuffer = allocator.buffer(validityBufferSize); - validityBuffer.readerIndex(0); + @Override + protected void allocateValidityBuffer(final long validityBufferSize) { + super.allocateValidityBuffer(validityBufferSize); refreshValueCapacity(); } @@ -359,7 +358,7 @@ public int getBufferSizeFor(final int count) { if (count == 0) { return 0; } - return (count * typeWidth) + getValidityBufferSizeFromCount(count); + return (count * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(count); } /** @@ -372,7 +371,7 @@ public int getBufferSize() { if (valueCount == 0) { return 0; } - return (valueCount * typeWidth) + getValidityBufferSizeFromCount(valueCount); + return (valueCount * typeWidth) + BitVectorHelper.getValidityBufferSizeFromCount(valueCount); } /** @@ -536,10 +535,10 @@ private void setReaderAndWriterIndex() { validityBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); if (typeWidth == 0) { /* specialized handling for BitVector */ - valueBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); + valueBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } else { valueBuffer.writerIndex((long) valueCount * typeWidth); } @@ -656,72 +655,18 @@ private void splitAndTransferValueBuffer( target.refreshValueCapacity(); } - /** - * Validity buffer has multiple cases of split and transfer depending on the starting position of - * the source index. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, BaseFixedWidthVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - /* slice */ - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); - target.refreshValueCapacity(); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } + @Override + protected void sliceAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { + final int firstByteSource = BitVectorHelper.byteIndex(startIndex); + final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); + + if (target.validityBuffer != null) { + target.validityBuffer.getReferenceManager().release(); } + ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); + target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); + ((BaseFixedWidthVector) target).refreshValueCapacity(); } /*----------------------------------------------------------------* diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java index 7e0d0affc6..3fac195786 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseLargeVariableWidthVector.java @@ -52,10 +52,8 @@ public abstract class BaseLargeVariableWidthVector extends BaseValueVector /* protected members */ public static final int OFFSET_WIDTH = 8; /* 8 byte unsigned int to track offsets */ protected static final byte[] emptyByteArray = new byte[] {}; - protected ArrowBuf validityBuffer; protected ArrowBuf valueBuffer; protected ArrowBuf offsetBuffer; - protected int valueCount; protected int lastSet; protected final Field field; @@ -375,14 +373,26 @@ private void setReaderAndWriterIndex() { valueBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); - offsetBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { final long lastDataOffset = getStartOffset(valueCount); - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); - offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); valueBuffer.writerIndex(lastDataOffset); } + // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`. + // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers + // in other libraries. According to Arrow spec, we should still output the offset buffer which + // is [0]. + final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH; + if (offsetBuffer.capacity() < requiredOffsetBufferSize) { + ArrowBuf newOffsetBuffer = allocateOffsetBuffer(requiredOffsetBufferSize); + if (offsetBuffer.capacity() > 0) { + newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); + } + offsetBuffer.getReferenceManager().release(); + offsetBuffer = newOffsetBuffer; + } + offsetBuffer.writerIndex(requiredOffsetBufferSize); } /** Same as {@link #allocateNewSafe()}. */ @@ -501,10 +511,9 @@ private ArrowBuf allocateOffsetBuffer(final long size) { } /* allocate validity buffer */ - private void allocateValidityBuffer(final long size) { - validityBuffer = allocator.buffer(size); - validityBuffer.readerIndex(0); - initValidityBuffer(); + @Override + protected void allocateValidityBuffer(final long size) { + super.allocateValidityBuffer(size); } /** @@ -633,7 +642,7 @@ public int getBufferSizeFor(final int valueCount) { return 0; } - final long validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final long validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); final long offsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH; /* get the end offset for this valueCount */ final long dataBufferSize = getStartOffset(valueCount); @@ -809,69 +818,17 @@ private void splitAndTransferOffsetBuffer( target.valueBuffer = transferBuffer(slicedBuffer, target.allocator); } - /* - * Transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, BaseLargeVariableWidthVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; + @Override + protected void sliceAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { + final int firstByteSource = BitVectorHelper.byteIndex(startIndex); + final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } + if (target.validityBuffer != null) { + target.validityBuffer.getReferenceManager().release(); } + target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); + target.validityBuffer.getReferenceManager().retain(); } /*----------------------------------------------------------------* diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java index 9befcb890f..37dfa20616 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java @@ -48,6 +48,10 @@ public abstract class BaseValueVector implements ValueVector { protected volatile FieldReader fieldReader; + protected ArrowBuf validityBuffer; + + protected int valueCount; + protected BaseValueVector(BufferAllocator allocator) { this.allocator = Preconditions.checkNotNull(allocator, "allocator cannot be null"); } @@ -110,7 +114,14 @@ protected ArrowBuf releaseBuffer(ArrowBuf buffer) { return buffer; } - /* number of bytes for the validity buffer for the given valueCount */ + /** + * Compute the size of validity buffer required to manage a given number of elements in a vector. + * + * @param valueCount number of elements in the vector + * @return buffer size + * @deprecated -- use {@link BitVectorHelper#getValidityBufferSizeFromCount} instead. + */ + @Deprecated(forRemoval = true, since = "18.4.0") protected static int getValidityBufferSizeFromCount(final int valueCount) { return DataSizeRoundingUtil.divideBy8Ceil(valueCount); } @@ -248,4 +259,116 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { throw new UnsupportedOperationException(); } + + /** + * Transfer the validity buffer from `validityBuffer` to the target vector's `validityBuffer`. + * Start at `startIndex` and copy `length` number of elements. If the starting index is 8 byte + * aligned, then the buffer is sliced from that index and ownership is transferred. If not, + * individual bytes are copied. + * + * @param startIndex starting index + * @param length number of elements to be copied + * @param target target vector + */ + protected void splitAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { + int offset = startIndex % 8; + + if (length <= 0) { + return; + } + if (offset == 0) { + sliceAndTransferValidityBuffer(startIndex, length, target); + } else { + copyValidityBuffer(startIndex, length, target); + } + } + + /** + * If the start index is 8 byte aligned, slice `validityBuffer` and transfer ownership to + * `target`'s `validityBuffer`. + * + * @param startIndex starting index + * @param length number of elements to be copied + * @param target target vector + */ + protected void sliceAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { + final int firstByteSource = BitVectorHelper.byteIndex(startIndex); + final int byteSizeTarget = getValidityBufferSizeFromCount(length); + + if (target.validityBuffer != null) { + target.validityBuffer.getReferenceManager().release(); + } + target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); + target.validityBuffer.getReferenceManager().retain(1); + } + + /** + * Allocate new validity buffer for `target` and copy bytes from `validityBuffer`. Precise details + * in the comments below. + * + * @param startIndex starting index + * @param length number of elements to be copied + * @param target target vector + */ + protected void copyValidityBuffer(int startIndex, int length, BaseValueVector target) { + final int firstByteSource = BitVectorHelper.byteIndex(startIndex); + final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); + final int byteSizeTarget = getValidityBufferSizeFromCount(length); + final int offset = startIndex % 8; + + /* Copy data + * When the first bit starts from the middle of a byte (offset != 0), + * copy data from src BitVector. + * Each byte in the target is composed by a part in i-th byte, + * another part in (i+1)-th byte. + */ + target.allocateValidityBuffer(byteSizeTarget); + + for (int i = 0; i < byteSizeTarget - 1; i++) { + byte b1 = + BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset); + byte b2 = + BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset); + + target.validityBuffer.setByte(i, (b1 + b2)); + } + + /* Copying the last piece is done in the following manner: + * if the source vector has 1 or more bytes remaining, we copy + * the last piece as a byte formed by shifting data + * from the current byte and the next byte. + * + * if the source vector has no more bytes remaining + * (we are at the last byte), we copy the last piece as a byte + * by shifting data from the current byte. + */ + if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { + byte b1 = + BitVectorHelper.getBitsFromCurrentByte( + this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); + byte b2 = + BitVectorHelper.getBitsFromNextByte( + this.validityBuffer, firstByteSource + byteSizeTarget, offset); + + target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); + } else { + byte b1 = + BitVectorHelper.getBitsFromCurrentByte( + this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); + target.validityBuffer.setByte(byteSizeTarget - 1, b1); + } + } + + /** + * Allocate new validity buffer for when the bytes need to be copied over. + * + * @param byteSizeTarget desired size of the buffer + */ + protected void allocateValidityBuffer(long byteSizeTarget) { + validityBuffer = allocator.buffer(byteSizeTarget); + validityBuffer.readerIndex(0); + validityBuffer.setZero(0, validityBuffer.capacity()); + } } diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java index 1609e64ca5..d5bd167256 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthVector.java @@ -50,10 +50,8 @@ public abstract class BaseVariableWidthVector extends BaseValueVector /* protected members */ public static final int OFFSET_WIDTH = 4; /* 4 byte unsigned int to track offsets */ protected static final byte[] emptyByteArray = new byte[] {}; - protected ArrowBuf validityBuffer; protected ArrowBuf valueBuffer; protected ArrowBuf offsetBuffer; - protected int valueCount; protected int lastSet; protected final Field field; @@ -87,7 +85,7 @@ public String getName() { /* TODO: * Once the entire hierarchy has been refactored, move common functions - * like getNullCount(), splitAndTransferValidityBuffer to top level + * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be @@ -391,14 +389,26 @@ private void setReaderAndWriterIndex() { valueBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); - offsetBuffer.writerIndex(0); valueBuffer.writerIndex(0); } else { final int lastDataOffset = getStartOffset(valueCount); - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); - offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); valueBuffer.writerIndex(lastDataOffset); } + // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`. + // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers + // in other libraries. According to Arrow spec, we should still output the offset buffer which + // is [0]. + final long requiredOffsetBufferSize = (long) (valueCount + 1) * OFFSET_WIDTH; + if (offsetBuffer.capacity() < requiredOffsetBufferSize) { + ArrowBuf newOffsetBuffer = allocateOffsetBuffer(requiredOffsetBufferSize); + if (offsetBuffer.capacity() > 0) { + newOffsetBuffer.setBytes(0, offsetBuffer, 0, offsetBuffer.capacity()); + } + offsetBuffer.getReferenceManager().release(); + offsetBuffer = newOffsetBuffer; + } + offsetBuffer.writerIndex(requiredOffsetBufferSize); } /** Same as {@link #allocateNewSafe()}. */ @@ -519,11 +529,9 @@ private ArrowBuf allocateOffsetBuffer(final long size) { } /* allocate validity buffer */ - private void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - initValidityBuffer(); + @Override + protected void allocateValidityBuffer(final long size) { + super.allocateValidityBuffer(size); } /** @@ -673,7 +681,7 @@ public int getBufferSizeFor(final int valueCount) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH; /* get the end offset for this valueCount */ final int dataBufferSize = offsetBuffer.getInt((long) valueCount * OFFSET_WIDTH); @@ -856,70 +864,17 @@ private void splitAndTransferOffsetBuffer( target.valueBuffer = transferBuffer(slicedBuffer, target.allocator); } - /* - * Transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, BaseVariableWidthVector target) { - if (length <= 0) { - return; - } - + @Override + protected void sliceAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { final int firstByteSource = BitVectorHelper.byteIndex(startIndex); - final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - final int byteSizeTarget = getValidityBufferSizeFromCount(length); - final int offset = startIndex % 8; - - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); - return; - } - - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); + final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); + if (target.validityBuffer != null) { + target.validityBuffer.getReferenceManager().release(); } + final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); + target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); } /*----------------------------------------------------------------* diff --git a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java index beda91dc3f..ea9de8320e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java @@ -78,13 +78,11 @@ public abstract class BaseVariableWidthViewVector extends BaseValueVector // The third 4 bytes of view are allocated for buffer index public static final int BUF_INDEX_WIDTH = 4; public static final byte[] EMPTY_BYTE_ARRAY = new byte[] {}; - protected ArrowBuf validityBuffer; // The view buffer is used to store the variable width view elements protected ArrowBuf viewBuffer; // The external buffer which stores the long strings protected List dataBuffers; protected int initialDataBufferSize; - protected int valueCount; protected int lastSet; protected final Field field; @@ -117,7 +115,7 @@ public String getName() { /* TODO: * Once the entire hierarchy has been refactored, move common functions - * like getNullCount(), splitAndTransferValidityBuffer to top level + * like getNullCount() to top level * base class BaseValueVector. * * Along with this, some class members (validityBuffer) can also be @@ -129,12 +127,6 @@ public String getName() { * the top class as of now is not a good idea. */ - /* TODO: - * Implement TransferPair functionality - * https://github.com/apache/arrow/issues/40932 - * - */ - /** * Get buffer that manages the validity (NULL or NON-NULL nature) of elements in the vector. * Consider it as a buffer for internal bit vector data structure. @@ -400,7 +392,7 @@ private void setReaderAndWriterIndex() { validityBuffer.writerIndex(0); viewBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); viewBuffer.writerIndex(valueCount * ELEMENT_SIZE); } } @@ -683,7 +675,7 @@ public int getBufferSizeFor(final int valueCount) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); final int viewBufferSize = valueCount * ELEMENT_SIZE; final int dataBufferSize = getDataBufferSize(); return validityBufferSize + viewBufferSize + dataBufferSize; @@ -854,77 +846,22 @@ public void splitAndTransferTo(int startIndex, int length, BaseVariableWidthView } /* allocate validity buffer */ - private void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - initValidityBuffer(); + @Override + protected void allocateValidityBuffer(final long size) { + super.allocateValidityBuffer(size); } - /* - * Transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, BaseVariableWidthViewVector target) { - if (length <= 0) { - return; - } - + @Override + protected void sliceAndTransferValidityBuffer( + int startIndex, int length, BaseValueVector target) { final int firstByteSource = BitVectorHelper.byteIndex(startIndex); - final int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - final int byteSizeTarget = getValidityBufferSizeFromCount(length); - final int offset = startIndex % 8; - - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); - return; - } + final int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(this.validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte(this.validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - this.validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - this.validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); + if (target.validityBuffer != null) { + target.validityBuffer.getReferenceManager().release(); } + final ArrowBuf slicedValidityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); + target.validityBuffer = transferBuffer(slicedValidityBuffer, target.allocator); } /** diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVector.java b/vector/src/main/java/org/apache/arrow/vector/BitVector.java index f8e3342625..ecee02f665 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BitVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/BitVector.java @@ -98,7 +98,7 @@ public MinorType getMinorType() { */ @Override public void setInitialCapacity(int valueCount) { - final int size = getValidityBufferSizeFromCount(valueCount); + final int size = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); if (size * 2L > MAX_ALLOCATION_SIZE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } @@ -121,7 +121,7 @@ public int getBufferSizeFor(final int count) { if (count == 0) { return 0; } - return 2 * getValidityBufferSizeFromCount(count); + return 2 * BitVectorHelper.getValidityBufferSizeFromCount(count); } /** @@ -165,7 +165,7 @@ private ArrowBuf splitAndTransferBuffer( int startIndex, int length, ArrowBuf sourceBuffer, ArrowBuf destBuffer) { int firstByteSource = BitVectorHelper.byteIndex(startIndex); int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); + int byteSizeTarget = BitVectorHelper.getValidityBufferSizeFromCount(length); int offset = startIndex % 8; if (length > 0) { diff --git a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java index 0ac56691a6..bc2c3da98f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java +++ b/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java @@ -135,11 +135,11 @@ public static void setValidityBit(ArrowBuf validityBuffer, int index, int value) public static ArrowBuf setValidityBit( ArrowBuf validityBuffer, BufferAllocator allocator, int valueCount, int index, int value) { if (validityBuffer == null) { - validityBuffer = allocator.buffer(getValidityBufferSize(valueCount)); + validityBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount)); } setValidityBit(validityBuffer, index, value); if (index == (valueCount - 1)) { - validityBuffer.writerIndex(getValidityBufferSize(valueCount)); + validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); } return validityBuffer; @@ -165,7 +165,7 @@ public static int get(final ArrowBuf buffer, int index) { * @param valueCount number of elements in the vector * @return buffer size */ - public static int getValidityBufferSize(int valueCount) { + public static int getValidityBufferSizeFromCount(int valueCount) { return DataSizeRoundingUtil.divideBy8Ceil(valueCount); } @@ -182,7 +182,7 @@ public static int getNullCount(final ArrowBuf validityBuffer, final int valueCou return 0; } int count = 0; - final int sizeInBytes = getValidityBufferSize(valueCount); + final int sizeInBytes = getValidityBufferSizeFromCount(valueCount); // If value count is not a multiple of 8, then calculate number of used bits in the last byte final int remainder = valueCount % 8; final int fullBytesCount = remainder == 0 ? sizeInBytes : sizeInBytes - 1; @@ -233,7 +233,7 @@ public static boolean checkAllBitsEqualTo( if (valueCount == 0) { return true; } - final int sizeInBytes = getValidityBufferSize(valueCount); + final int sizeInBytes = getValidityBufferSizeFromCount(valueCount); // boundary check validityBuffer.checkBytes(0, sizeInBytes); @@ -325,7 +325,7 @@ public static ArrowBuf loadValidityBuffer( sourceValidityBuffer == null || sourceValidityBuffer.capacity() == 0; if (isValidityBufferNull && (fieldNode.getNullCount() == 0 || fieldNode.getNullCount() == valueCount)) { - newBuffer = allocator.buffer(getValidityBufferSize(valueCount)); + newBuffer = allocator.buffer(getValidityBufferSizeFromCount(valueCount)); newBuffer.setZero(0, newBuffer.capacity()); if (fieldNode.getNullCount() != 0) { /* all NULLs */ diff --git a/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java b/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java index e22a973f3b..61a5574898 100644 --- a/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/FixedWidthVector.java @@ -31,4 +31,7 @@ public interface FixedWidthVector extends ElementAddressableVector { /** Zero out the underlying buffer backing this vector. */ void zeroVector(); + + /** Get the width of the type in bytes. */ + int getTypeWidth(); } diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java index abaefcfc12..50f2f066cc 100644 --- a/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampMicroTZVector.java @@ -155,12 +155,13 @@ public void set(int index, NullableTimeStampMicroTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); - } else if (!this.timeZone.equals(holder.timezone)) { - throw new IllegalArgumentException( - String.format( - "holder.timezone: %s not equal to vector timezone: %s", - holder.timezone, this.timeZone)); } else if (holder.isSet > 0) { + if (!this.timeZone.equals(holder.timezone)) { + throw new IllegalArgumentException( + String.format( + "holder.timezone: %s not equal to vector timezone: %s", + holder.timezone, this.timeZone)); + } BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java index b5e5fb1be1..9e4998396c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampMilliTZVector.java @@ -155,12 +155,13 @@ public void set(int index, NullableTimeStampMilliTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); - } else if (!this.timeZone.equals(holder.timezone)) { - throw new IllegalArgumentException( - String.format( - "holder.timezone: %s not equal to vector timezone: %s", - holder.timezone, this.timeZone)); } else if (holder.isSet > 0) { + if (!this.timeZone.equals(holder.timezone)) { + throw new IllegalArgumentException( + String.format( + "holder.timezone: %s not equal to vector timezone: %s", + holder.timezone, this.timeZone)); + } BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java index 2386b3a859..b44b3da8d3 100644 --- a/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampNanoTZVector.java @@ -154,12 +154,13 @@ public Long getObject(int index) { public void set(int index, NullableTimeStampNanoTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); - } else if (!this.timeZone.equals(holder.timezone)) { - throw new IllegalArgumentException( - String.format( - "holder.timezone: %s not equal to vector timezone: %s", - holder.timezone, this.timeZone)); } else if (holder.isSet > 0) { + if (!this.timeZone.equals(holder.timezone)) { + throw new IllegalArgumentException( + String.format( + "holder.timezone: %s not equal to vector timezone: %s", + holder.timezone, this.timeZone)); + } BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { diff --git a/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java b/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java index f1774f2703..a64a87f699 100644 --- a/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/TimeStampSecTZVector.java @@ -150,12 +150,13 @@ public Long getObject(int index) { public void set(int index, NullableTimeStampSecTZHolder holder) throws IllegalArgumentException { if (holder.isSet < 0) { throw new IllegalArgumentException(); - } else if (!this.timeZone.equals(holder.timezone)) { - throw new IllegalArgumentException( - String.format( - "holder.timezone: %s not equal to vector timezone: %s", - holder.timezone, this.timeZone)); } else if (holder.isSet > 0) { + if (!this.timeZone.equals(holder.timezone)) { + throw new IllegalArgumentException( + String.format( + "holder.timezone: %s not equal to vector timezone: %s", + holder.timezone, this.timeZone)); + } BitVectorHelper.setBit(validityBuffer, index); setValue(index, holder.value); } else { diff --git a/vector/src/main/java/org/apache/arrow/vector/UuidVector.java b/vector/src/main/java/org/apache/arrow/vector/UuidVector.java new file mode 100644 index 0000000000..e1e61a5a2e --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/UuidVector.java @@ -0,0 +1,458 @@ +/* + * 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.arrow.vector; + +import static org.apache.arrow.vector.extension.UuidType.UUID_BYTE_WIDTH; + +import java.nio.ByteBuffer; +import java.util.UUID; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.ArrowBufPointer; +import org.apache.arrow.memory.util.ByteFunctionHelpers; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.complex.impl.UuidReaderImpl; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.apache.arrow.vector.holders.UuidHolder; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.CallBack; +import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.UuidUtility; + +/** + * Vector implementation for UUID values using {@link UuidType}. + * + *

Supports setting and retrieving UUIDs with efficient storage and nullable value handling. + * + *

Usage: + * + * 

{@code
+ * UuidVector vector = new UuidVector("uuid_col", allocator);
+ * vector.set(0, UUID.randomUUID());
+ * UUID value = vector.getObject(0);
+ * }
+ * + * @see UuidType + * @see UuidHolder + * @see NullableUuidHolder + */ +public class UuidVector extends ExtensionTypeVector + implements ValueIterableVector, FixedWidthVector { + private final Field field; + + /** The fixed byte width of UUID values (16 bytes). */ + public static final int TYPE_WIDTH = UUID_BYTE_WIDTH; + + /** + * Constructs a UUID vector with the given name, allocator, and underlying vector. + * + * @param name the name of the vector + * @param allocator the buffer allocator + * @param underlyingVector the underlying FixedSizeBinaryVector for storage + */ + public UuidVector( + String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) { + super(name, allocator, underlyingVector); + this.field = new Field(name, FieldType.nullable(UuidType.INSTANCE), null); + } + + /** + * Constructs a UUID vector with the given name, field type, allocator, and underlying vector. + * + * @param name the name of the vector + * @param fieldType the field type (should contain UuidType) + * @param allocator the buffer allocator + * @param underlyingVector the underlying FixedSizeBinaryVector for storage + */ + public UuidVector( + String name, + FieldType fieldType, + BufferAllocator allocator, + FixedSizeBinaryVector underlyingVector) { + super(name, allocator, underlyingVector); + this.field = new Field(name, fieldType, null); + } + + /** + * Constructs a UUID vector with the given name and allocator. + * + *

Creates a new underlying FixedSizeBinaryVector with 16-byte width. + * + * @param name the name of the vector + * @param allocator the buffer allocator + */ + public UuidVector(String name, BufferAllocator allocator) { + super(name, allocator, new FixedSizeBinaryVector(name, allocator, UUID_BYTE_WIDTH)); + this.field = new Field(name, FieldType.nullable(UuidType.INSTANCE), null); + } + + /** + * Constructs a UUID vector from a field and allocator. + * + * @param field the field definition (should contain UuidType) + * @param allocator the buffer allocator + */ + public UuidVector(Field field, BufferAllocator allocator) { + super( + field.getName(), + allocator, + new FixedSizeBinaryVector(field.getName(), allocator, UUID_BYTE_WIDTH)); + this.field = field; + } + + @Override + public UUID getObject(int index) { + if (isSet(index) == 0) { + return null; + } + final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index)); + return new UUID(bb.getLong(), bb.getLong()); + } + + @Override + public int hashCode(int index) { + return hashCode(index, null); + } + + @Override + public int hashCode(int index, ArrowBufHasher hasher) { + int start = this.getStartOffset(index); + return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, start + UUID_BYTE_WIDTH); + } + + /** + * Checks if the value at the given index is set (non-null). + * + * @param index the index to check + * @return 1 if the value is set, 0 if null + */ + public int isSet(int index) { + return getUnderlyingVector().isSet(index); + } + + /** + * Reads the UUID value at the given index into a NullableUuidHolder. + * + * @param index the index to read from + * @param holder the holder to populate with the UUID data + */ + public void get(int index, NullableUuidHolder holder) { + Preconditions.checkArgument(index >= 0, "Cannot get negative index in UUID vector."); + if (isSet(index) == 0) { + holder.isSet = 0; + return; + } + holder.isSet = 1; + holder.buffer = getDataBuffer(); + holder.start = getStartOffset(index); + } + + /** + * Calculates the byte offset for a given index in the data buffer. + * + * @param index the index of the UUID value + * @return the byte offset in the data buffer + */ + public final int getStartOffset(int index) { + return index * UUID_BYTE_WIDTH; + } + + /** + * Sets the UUID value at the given index. + * + * @param index the index to set + * @param value the UUID value to set, or null to set a null value + */ + public void set(int index, UUID value) { + if (value != null) { + set(index, UuidUtility.getBytesFromUUID(value)); + } else { + getUnderlyingVector().setNull(index); + } + } + + /** + * Sets the UUID value at the given index from a UuidHolder. + * + * @param index the index to set + * @param holder the holder containing the UUID data + */ + public void set(int index, UuidHolder holder) { + this.set(index, holder.buffer, holder.start); + } + + /** + * Sets the UUID value at the given index from a NullableUuidHolder. + * + * @param index the index to set + * @param holder the holder containing the UUID data + */ + public void set(int index, NullableUuidHolder holder) { + if (holder.isSet == 0) { + getUnderlyingVector().setNull(index); + } else { + this.set(index, holder.buffer, holder.start); + } + } + + /** + * Sets the UUID value at the given index by copying from a source buffer. + * + * @param index the index to set + * @param source the source buffer to copy from + * @param sourceOffset the offset in the source buffer where the UUID data starts + */ + public void set(int index, ArrowBuf source, int sourceOffset) { + Preconditions.checkNotNull(source, "Cannot set UUID vector, the source buffer is null."); + + BitVectorHelper.setBit(getUnderlyingVector().getValidityBuffer(), index); + getUnderlyingVector() + .getDataBuffer() + .setBytes((long) index * UUID_BYTE_WIDTH, source, sourceOffset, UUID_BYTE_WIDTH); + } + + /** + * Sets the UUID value at the given index from a byte array. + * + * @param index the index to set + * @param value the 16-byte array containing the UUID data + */ + public void set(int index, byte[] value) { + getUnderlyingVector().set(index, value); + } + + /** + * Sets the UUID value at the given index, expanding capacity if needed. + * + * @param index the index to set + * @param value the UUID value to set, or null to set a null value + */ + public void setSafe(int index, UUID value) { + if (value != null) { + setSafe(index, UuidUtility.getBytesFromUUID(value)); + } else { + getUnderlyingVector().setNull(index); + } + } + + /** + * Sets the UUID value at the given index from a NullableUuidHolder, expanding capacity if needed. + * + * @param index the index to set + * @param holder the holder containing the UUID data, or null to set a null value + */ + public void setSafe(int index, NullableUuidHolder holder) { + if (holder == null || holder.isSet == 0) { + getUnderlyingVector().setNull(index); + } else { + this.setSafe(index, holder.buffer, holder.start); + } + } + + /** + * Sets the UUID value at the given index from a UuidHolder, expanding capacity if needed. + * + * @param index the index to set + * @param holder the holder containing the UUID data + */ + public void setSafe(int index, UuidHolder holder) { + this.setSafe(index, holder.buffer, holder.start); + } + + /** + * Sets the UUID value at the given index by copying from a source buffer, expanding capacity if + * needed. + * + * @param index the index to set + * @param buffer the source buffer to copy from + * @param start the offset in the source buffer where the UUID data starts + */ + public void setSafe(int index, ArrowBuf buffer, int start) { + getUnderlyingVector().handleSafe(index); + this.set(index, buffer, start); + } + + /** + * Sets the UUID value at the given index from a byte array, expanding capacity if needed. + * + * @param index the index to set + * @param value the 16-byte array containing the UUID data + */ + public void setSafe(int index, byte[] value) { + getUnderlyingVector().setIndexDefined(index); + getUnderlyingVector().setSafe(index, value); + } + + /** + * Sets the UUID value at the given index from an ArrowBuf, expanding capacity if needed. + * + * @param index the index to set + * @param value the buffer containing the 16-byte UUID data + */ + public void setSafe(int index, ArrowBuf value) { + getUnderlyingVector().setSafe(index, value); + } + + @Override + public void copyFrom(int fromIndex, int thisIndex, ValueVector from) { + getUnderlyingVector() + .copyFromSafe(fromIndex, thisIndex, ((UuidVector) from).getUnderlyingVector()); + } + + @Override + public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { + getUnderlyingVector() + .copyFromSafe(fromIndex, thisIndex, ((UuidVector) from).getUnderlyingVector()); + } + + @Override + public Field getField() { + return field; + } + + @Override + public ArrowBufPointer getDataPointer(int i) { + return getUnderlyingVector().getDataPointer(i); + } + + @Override + public ArrowBufPointer getDataPointer(int i, ArrowBufPointer arrowBufPointer) { + return getUnderlyingVector().getDataPointer(i, arrowBufPointer); + } + + @Override + public void allocateNew(int valueCount) { + getUnderlyingVector().allocateNew(valueCount); + } + + @Override + public void zeroVector() { + getUnderlyingVector().zeroVector(); + } + + @Override + public TransferPair makeTransferPair(ValueVector to) { + return new TransferImpl((UuidVector) to); + } + + @Override + protected FieldReader getReaderImpl() { + return new UuidReaderImpl(this); + } + + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator) { + return new TransferImpl(field, allocator); + } + + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(field, allocator); + } + + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator) { + return new TransferImpl(ref, allocator); + } + + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator, CallBack callBack) { + return getTransferPair(ref, allocator); + } + + @Override + public TransferPair getTransferPair(BufferAllocator allocator) { + return getTransferPair(this.getField().getName(), allocator); + } + + @Override + public int getTypeWidth() { + return UUID_BYTE_WIDTH; + } + + /** {@link TransferPair} for {@link UuidVector}. */ + public class TransferImpl implements TransferPair { + UuidVector to; + + /** + * Constructs a transfer pair with the given target vector. + * + * @param to the target UUID vector + */ + public TransferImpl(UuidVector to) { + this.to = to; + } + + /** + * Constructs a transfer pair, creating a new target vector from the field and allocator. + * + * @param field the field definition for the target vector + * @param allocator the buffer allocator for the target vector + */ + public TransferImpl(Field field, BufferAllocator allocator) { + this.to = new UuidVector(field, allocator); + } + + /** + * Constructs a transfer pair, creating a new target vector with the given name and allocator. + * + * @param ref the name for the target vector + * @param allocator the buffer allocator for the target vector + */ + public TransferImpl(String ref, BufferAllocator allocator) { + this.to = new UuidVector(ref, allocator); + } + + /** + * Gets the target vector of this transfer pair. + * + * @return the target UUID vector + */ + public UuidVector getTo() { + return this.to; + } + + /** Transfers ownership of data from the source vector to the target vector. */ + public void transfer() { + getUnderlyingVector().transferTo(to.getUnderlyingVector()); + } + + /** + * Splits and transfers a range of values from the source vector to the target vector. + * + * @param startIndex the starting index in the source vector + * @param length the number of values to transfer + */ + public void splitAndTransfer(int startIndex, int length) { + getUnderlyingVector().splitAndTransferTo(startIndex, length, to.getUnderlyingVector()); + } + + /** + * Copies a value from the source vector to the target vector, expanding capacity if needed. + * + * @param fromIndex the index in the source vector + * @param toIndex the index in the target vector + */ + public void copyValueSafe(int fromIndex, int toIndex) { + to.copyFromSafe(fromIndex, toIndex, (ValueVector) UuidVector.this); + } + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java index a7cb9ced72..4c1fbf761a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java @@ -199,13 +199,18 @@ public FieldVector getVector(int index) { */ public VectorSchemaRoot addVector(int index, FieldVector vector) { Preconditions.checkNotNull(vector); - Preconditions.checkArgument(index >= 0 && index < fieldVectors.size()); + Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size()); List newVectors = new ArrayList<>(); - for (int i = 0; i < fieldVectors.size(); i++) { - if (i == index) { - newVectors.add(vector); + if (index == fieldVectors.size()) { + newVectors.addAll(fieldVectors); + newVectors.add(vector); + } else { + for (int i = 0; i < fieldVectors.size(); i++) { + if (i == index) { + newVectors.add(vector); + } + newVectors.add(fieldVectors.get(i)); } - newVectors.add(fieldVectors.get(i)); } return new VectorSchemaRoot(newVectors); } diff --git a/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java b/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java index abcf312c5e..bc2e3a6aab 100644 --- a/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/compare/RangeEqualsVisitor.java @@ -43,6 +43,7 @@ import org.apache.arrow.vector.complex.ListViewVector; import org.apache.arrow.vector.complex.NonNullableStructVector; import org.apache.arrow.vector.complex.RunEndEncodedVector; +import org.apache.arrow.vector.complex.RunEndEncodedVector.RangeIterator; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; @@ -270,42 +271,35 @@ protected boolean compareRunEndEncodedVectors(Range range) { RunEndEncodedVector leftVector = (RunEndEncodedVector) left; RunEndEncodedVector rightVector = (RunEndEncodedVector) right; - final int leftRangeEnd = range.getLeftStart() + range.getLength(); - final int rightRangeEnd = range.getRightStart() + range.getLength(); + final RunEndEncodedVector.RangeIterator leftIterator = + new RunEndEncodedVector.RangeIterator(leftVector, range.getLeftStart(), range.getLength()); + final RunEndEncodedVector.RangeIterator rightIterator = + new RunEndEncodedVector.RangeIterator( + rightVector, range.getRightStart(), range.getLength()); FieldVector leftValuesVector = leftVector.getValuesVector(); FieldVector rightValuesVector = rightVector.getValuesVector(); RangeEqualsVisitor innerVisitor = createInnerVisitor(leftValuesVector, rightValuesVector, null); - int leftLogicalIndex = range.getLeftStart(); - int rightLogicalIndex = range.getRightStart(); + while (nextRun(leftIterator, rightIterator)) { + int leftPhysicalIndex = leftIterator.getRunIndex(); + int rightPhysicalIndex = rightIterator.getRunIndex(); - while (leftLogicalIndex < leftRangeEnd) { - // TODO: implement it more efficient - // https://github.com/apache/arrow/issues/44157 - int leftPhysicalIndex = leftVector.getPhysicalIndex(leftLogicalIndex); - int rightPhysicalIndex = rightVector.getPhysicalIndex(rightLogicalIndex); - if (leftValuesVector.accept( - innerVisitor, new Range(leftPhysicalIndex, rightPhysicalIndex, 1))) { - int leftRunEnd = leftVector.getRunEnd(leftLogicalIndex); - int rightRunEnd = rightVector.getRunEnd(rightLogicalIndex); - - int leftRunLength = Math.min(leftRunEnd, leftRangeEnd) - leftLogicalIndex; - int rightRunLength = Math.min(rightRunEnd, rightRangeEnd) - rightLogicalIndex; - - if (leftRunLength != rightRunLength) { - return false; - } else { - leftLogicalIndex = leftRunEnd; - rightLogicalIndex = rightRunEnd; - } - } else { + if (leftIterator.getRunLength() != rightIterator.getRunLength() + || !leftValuesVector.accept( + innerVisitor, new Range(leftPhysicalIndex, rightPhysicalIndex, 1))) { return false; } } - return true; + return leftIterator.isEnd() && rightIterator.isEnd(); + } + + private static boolean nextRun(RangeIterator leftIterator, RangeIterator rightIterator) { + boolean left = leftIterator.nextRun(); + boolean right = rightIterator.nextRun(); + return left && right; } protected RangeEqualsVisitor createInnerVisitor( diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java index 2921e43cb6..a57fbe473f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/AbstractStructVector.java @@ -46,11 +46,13 @@ public abstract class AbstractStructVector extends AbstractContainerVector { private ConflictPolicy conflictPolicy; static { - String conflictPolicyStr = - System.getProperty(STRUCT_CONFLICT_POLICY_JVM, ConflictPolicy.CONFLICT_REPLACE.toString()); + String conflictPolicyStr = System.getProperty(STRUCT_CONFLICT_POLICY_JVM); if (conflictPolicyStr == null) { conflictPolicyStr = System.getenv(STRUCT_CONFLICT_POLICY_ENV); } + if (conflictPolicyStr == null) { + conflictPolicyStr = ConflictPolicy.CONFLICT_REPLACE.toString(); + } ConflictPolicy conflictPolicy; try { conflictPolicy = ConflictPolicy.valueOf(conflictPolicyStr.toUpperCase(Locale.ROOT)); @@ -62,11 +64,11 @@ public abstract class AbstractStructVector extends AbstractContainerVector { /** Policy to determine how to react when duplicate columns are encountered. */ public enum ConflictPolicy { - // Ignore the conflict and append the field. This is the default behaviour + // Ignore the conflict and append the field. CONFLICT_APPEND, // Keep the existing field and ignore the newer one. CONFLICT_IGNORE, - // Replace the existing field with the newer one. + // Replace the existing field with the newer one. This is the default behaviour CONFLICT_REPLACE, // Refuse the new field and error out. CONFLICT_ERROR diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java index 12edd6557b..fac3f86bba 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseLargeRepeatedValueViewVector.java @@ -52,7 +52,6 @@ public abstract class BaseLargeRepeatedValueViewVector extends BaseValueVector protected ArrowBuf sizeBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; - protected int valueCount; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH; private final String name; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java index fbe83bad52..ee1d65d3e3 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueVector.java @@ -54,7 +54,6 @@ public abstract class BaseRepeatedValueVector extends BaseValueVector protected ArrowBuf offsetBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; - protected int valueCount; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; private final String name; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java index e6213316b5..fd7a4ff2c6 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/BaseRepeatedValueViewVector.java @@ -52,7 +52,6 @@ public abstract class BaseRepeatedValueViewVector extends BaseValueVector protected ArrowBuf sizeBuffer; protected FieldVector vector; protected final CallBack repeatedCallBack; - protected int valueCount; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected long sizeAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * SIZE_WIDTH; private final String name; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java index c762eb5172..e3b4ab477f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java @@ -69,12 +69,10 @@ public static FixedSizeListVector empty(String name, int size, BufferAllocator a } private FieldVector vector; - private ArrowBuf validityBuffer; private final int listSize; private Field field; private UnionFixedSizeListReader reader; - private int valueCount; private int validityAllocationSizeInBytes; /** @@ -110,7 +108,8 @@ public FixedSizeListVector( this.listSize = ((ArrowType.FixedSizeList) field.getFieldType().getType()).getListSize(); Preconditions.checkArgument(listSize >= 0, "list size must be non-negative"); this.valueCount = 0; - this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); + this.validityAllocationSizeInBytes = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override @@ -189,7 +188,7 @@ public List getFieldBuffers() { private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } /** @@ -247,12 +246,10 @@ public boolean allocateNewSafe() { return success; } - private void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - validityAllocationSizeInBytes = curSize; - validityBuffer.setZero(0, validityBuffer.capacity()); + @Override + protected void allocateValidityBuffer(final long size) { + super.allocateValidityBuffer(size); + validityAllocationSizeInBytes = (int) size; } @Override @@ -268,7 +265,8 @@ private void reallocValidityBuffer() { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { - newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; + newAllocationSize = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } @@ -311,7 +309,7 @@ public UnionFixedSizeListWriter getWriter() { @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); vector.setInitialCapacity(numRecords * listSize); } @@ -328,7 +326,7 @@ public int getBufferSize() { if (getValueCount() == 0) { return 0; } - return getValidityBufferSizeFromCount(valueCount) + vector.getBufferSize(); + return BitVectorHelper.getValidityBufferSizeFromCount(valueCount) + vector.getBufferSize(); } @Override @@ -336,7 +334,7 @@ public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } - return getValidityBufferSizeFromCount(valueCount) + return BitVectorHelper.getValidityBufferSizeFromCount(valueCount) + vector.getBufferSizeFor(valueCount * listSize); } @@ -647,71 +645,6 @@ public void splitAndTransfer(int startIndex, int length) { to.setValueCount(length); } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, FixedSizeListVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java index ed075352c9..92dd3eaef7 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListVector.java @@ -31,6 +31,7 @@ import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.ByteFunctionHelpers; import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.memory.util.hash.ArrowBufHasher; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; @@ -94,11 +95,9 @@ public static LargeListVector empty(String name, BufferAllocator allocator) { protected ArrowBuf offsetBuffer; protected FieldVector vector; protected final CallBack callBack; - protected int valueCount; protected long offsetAllocationSizeInBytes = INITIAL_VALUE_ALLOCATION * OFFSET_WIDTH; protected String defaultDataVectorName = DATA_VECTOR_NAME; - protected ArrowBuf validityBuffer; protected UnionLargeListReader reader; private Field field; private int validityAllocationSizeInBytes; @@ -131,7 +130,8 @@ public LargeListVector(Field field, BufferAllocator allocator, CallBack callBack this.field = field; this.validityBuffer = allocator.getEmpty(); this.callBack = callBack; - this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); + this.validityAllocationSizeInBytes = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); this.lastSet = -1; this.offsetBuffer = allocator.getEmpty(); this.vector = vector == null ? DEFAULT_DATA_VECTOR : vector; @@ -156,7 +156,7 @@ public void initializeChildrenFromFields(List children) { @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); offsetAllocationSizeInBytes = (long) (numRecords + 1) * OFFSET_WIDTH; if (vector instanceof BaseFixedWidthVector || vector instanceof BaseVariableWidthVector) { vector.setInitialCapacity(numRecords * RepeatedValueVector.DEFAULT_REPEAT_PER_RECORD); @@ -184,7 +184,7 @@ public void setInitialCapacity(int numRecords) { */ @Override public void setInitialCapacity(int numRecords, double density) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); if ((numRecords * density) >= Integer.MAX_VALUE) { throw new OversizedAllocationException("Requested amount of memory is more than max allowed"); } @@ -309,11 +309,14 @@ private void setReaderAndWriterIndex() { offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); - offsetBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); - offsetBuffer.writerIndex((valueCount + 1) * OFFSET_WIDTH); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } + // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`. + // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers + // in other libraries. According to Arrow spec, we should still output the offset buffer which + // is [0]. + offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); } /** @@ -374,12 +377,10 @@ public boolean allocateNewSafe() { return success; } - private void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - validityAllocationSizeInBytes = curSize; - validityBuffer.setZero(0, validityBuffer.capacity()); + @Override + protected void allocateValidityBuffer(final long size) { + super.allocateValidityBuffer(size); + validityAllocationSizeInBytes = (int) size; } protected ArrowBuf allocateOffsetBuffer(final long size) { @@ -442,7 +443,8 @@ private void reallocValidityBuffer() { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { - newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; + newAllocationSize = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); @@ -692,71 +694,6 @@ public void splitAndTransfer(int startIndex, int length) { to.setValueCount(length); } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, LargeListVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; @@ -821,7 +758,7 @@ public int getBufferSize() { return 0; } final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH; - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + validityBufferSize + vector.getBufferSize(); } @@ -830,7 +767,7 @@ public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); long innerVectorValueCount = offsetBuffer.getLong((long) valueCount * OFFSET_WIDTH); return ((valueCount + 1) * OFFSET_WIDTH) @@ -928,10 +865,11 @@ public List getObject(int index) { if (isSet(index) == 0) { return null; } - final List vals = new JsonStringArrayList<>(); final long start = offsetBuffer.getLong((long) index * OFFSET_WIDTH); final long end = offsetBuffer.getLong(((long) index + 1L) * OFFSET_WIDTH); final ValueVector vv = getDataVector(); + final List vals = + new JsonStringArrayList<>(LargeMemoryUtil.checkedCastToInt(end - start)); for (long i = start; i < end; i++) { vals.add(vv.getObject(checkedCastToInt(i))); } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java index 84c6f03edb..2da7eb057e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/LargeListViewVector.java @@ -77,7 +77,6 @@ public class LargeListViewVector extends BaseLargeRepeatedValueViewVector implements PromotableVector, ValueIterableVector> { - protected ArrowBuf validityBuffer; protected UnionLargeListViewReader reader; private CallBack callBack; protected Field field; @@ -113,7 +112,8 @@ public LargeListViewVector(Field field, BufferAllocator allocator, CallBack call this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; - this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); + this.validityAllocationSizeInBytes = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override @@ -134,7 +134,7 @@ public void initializeChildrenFromFields(List children) { @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } @@ -157,7 +157,7 @@ public void setInitialCapacity(int numRecords) { */ @Override public void setInitialCapacity(int numRecords, double density) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } @@ -176,7 +176,7 @@ public void setInitialCapacity(int numRecords, double density) { */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } @@ -226,7 +226,7 @@ private void setReaderAndWriterIndex() { offsetBuffer.writerIndex(0); sizeBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH); sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH); } @@ -284,12 +284,10 @@ public boolean allocateNewSafe() { return success; } + @Override protected void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - validityAllocationSizeInBytes = curSize; - validityBuffer.setZero(0, validityBuffer.capacity()); + super.allocateValidityBuffer(size); + validityAllocationSizeInBytes = (int) size; } @Override @@ -323,7 +321,8 @@ private long getNewAllocationSize(int currentBufferCapacity) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { - newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; + newAllocationSize = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); @@ -529,71 +528,6 @@ public void splitAndTransfer(int startIndex, int length) { } } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer( - int startIndex, int length, LargeListViewVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; @@ -629,7 +563,7 @@ public int getBufferSize() { } final int offsetBufferSize = valueCount * OFFSET_WIDTH; final int sizeBufferSize = valueCount * SIZE_WIDTH; - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize(); } @@ -644,7 +578,7 @@ public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } @@ -738,10 +672,10 @@ public List getObject(int index) { if (isSet(index) == 0) { return null; } - final List vals = new JsonStringArrayList<>(); final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); final int end = start + sizeBuffer.getInt((index) * SIZE_WIDTH); final ValueVector vv = getDataVector(); + final List vals = new JsonStringArrayList<>(end - start); for (int i = start; i < end; i++) { vals.add(vv.getObject(checkedCastToInt(i))); } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java index 3daeb6d77b..6c3993df63 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java @@ -74,7 +74,6 @@ public static ListVector empty(String name, BufferAllocator allocator) { return new ListVector(name, allocator, FieldType.nullable(ArrowType.List.INSTANCE), null); } - protected ArrowBuf validityBuffer; protected UnionListReader reader; private CallBack callBack; protected Field field; @@ -108,7 +107,8 @@ public ListVector(Field field, BufferAllocator allocator, CallBack callBack) { this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; - this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); + this.validityAllocationSizeInBytes = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); this.lastSet = -1; } @@ -130,7 +130,7 @@ public void initializeChildrenFromFields(List children) { @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } @@ -153,7 +153,7 @@ public void setInitialCapacity(int numRecords) { */ @Override public void setInitialCapacity(int numRecords, double density) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } @@ -172,7 +172,7 @@ public void setInitialCapacity(int numRecords, double density) { */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } @@ -267,11 +267,14 @@ private void setReaderAndWriterIndex() { offsetBuffer.readerIndex(0); if (valueCount == 0) { validityBuffer.writerIndex(0); - offsetBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); - offsetBuffer.writerIndex((valueCount + 1) * OFFSET_WIDTH); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); } + // IPC serializer will determine readable bytes based on `readerIndex` and `writerIndex`. + // Both are set to 0 means 0 bytes are written to the IPC stream which will crash IPC readers + // in other libraries. According to Arrow spec, we should still output the offset buffer which + // is [0]. + offsetBuffer.writerIndex((long) (valueCount + 1) * OFFSET_WIDTH); } /** @@ -323,12 +326,10 @@ public boolean allocateNewSafe() { return success; } + @Override protected void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - validityAllocationSizeInBytes = curSize; - validityBuffer.setZero(0, validityBuffer.capacity()); + super.allocateValidityBuffer(size); + validityAllocationSizeInBytes = (int) size; } /** @@ -366,7 +367,8 @@ private long getNewAllocationSize(int currentBufferCapacity) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { - newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; + newAllocationSize = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); @@ -573,70 +575,6 @@ public void splitAndTransfer(int startIndex, int length) { } } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer(int startIndex, int length, ListVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; @@ -678,7 +616,7 @@ public int getBufferSize() { return 0; } final int offsetBufferSize = (valueCount + 1) * OFFSET_WIDTH; - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + validityBufferSize + vector.getBufferSize(); } @@ -687,7 +625,7 @@ public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } @@ -784,10 +722,10 @@ public List getObject(int index) { if (isSet(index) == 0) { return null; } - final List vals = new JsonStringArrayList<>(); final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); final int end = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); final ValueVector vv = getDataVector(); + final List vals = new JsonStringArrayList<>(end - start); for (int i = start; i < end; i++) { vals.add(vv.getObject(i)); } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java index 9b4e6b4c0c..d41f61e291 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/ListViewVector.java @@ -76,7 +76,6 @@ public class ListViewVector extends BaseRepeatedValueViewVector implements PromotableVector, ValueIterableVector> { - protected ArrowBuf validityBuffer; protected UnionListViewReader reader; private CallBack callBack; protected Field field; @@ -112,7 +111,8 @@ public ListViewVector(Field field, BufferAllocator allocator, CallBack callBack) this.validityBuffer = allocator.getEmpty(); this.field = field; this.callBack = callBack; - this.validityAllocationSizeInBytes = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); + this.validityAllocationSizeInBytes = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION); } @Override @@ -133,7 +133,7 @@ public void initializeChildrenFromFields(List children) { @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } @@ -156,7 +156,7 @@ public void setInitialCapacity(int numRecords) { */ @Override public void setInitialCapacity(int numRecords, double density) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } @@ -175,7 +175,7 @@ public void setInitialCapacity(int numRecords, double density) { */ @Override public void setInitialTotalCapacity(int numRecords, int totalNumberOfElements) { - validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); + validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSizeFromCount(numRecords); super.setInitialTotalCapacity(numRecords, totalNumberOfElements); } @@ -225,9 +225,9 @@ private void setReaderAndWriterIndex() { offsetBuffer.writerIndex(0); sizeBuffer.writerIndex(0); } else { - validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); - offsetBuffer.writerIndex(valueCount * OFFSET_WIDTH); - sizeBuffer.writerIndex(valueCount * SIZE_WIDTH); + validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSizeFromCount(valueCount)); + offsetBuffer.writerIndex((long) valueCount * OFFSET_WIDTH); + sizeBuffer.writerIndex((long) valueCount * SIZE_WIDTH); } } @@ -283,12 +283,10 @@ public boolean allocateNewSafe() { return success; } + @Override protected void allocateValidityBuffer(final long size) { - final int curSize = (int) size; - validityBuffer = allocator.buffer(curSize); - validityBuffer.readerIndex(0); - validityAllocationSizeInBytes = curSize; - validityBuffer.setZero(0, validityBuffer.capacity()); + super.allocateValidityBuffer(size); + validityAllocationSizeInBytes = (int) size; } @Override @@ -322,7 +320,8 @@ private long getNewAllocationSize(int currentBufferCapacity) { if (validityAllocationSizeInBytes > 0) { newAllocationSize = validityAllocationSizeInBytes; } else { - newAllocationSize = getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; + newAllocationSize = + BitVectorHelper.getValidityBufferSizeFromCount(INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); @@ -446,14 +445,22 @@ public int hashCode(int index, ArrowBufHasher hasher) { return ArrowBufPointer.NULL_HASH_CODE; } int hash = 0; - final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); - final int end = sizeBuffer.getInt(index * OFFSET_WIDTH); + final int start = getElementStartIndex(index); + final int end = getElementEndIndex(index); for (int i = start; i < end; i++) { hash = ByteFunctionHelpers.combineHash(hash, vector.hashCode(i, hasher)); } return hash; } + private void setElementOffsetBuffer(int index, int value) { + offsetBuffer.setInt((long) index * OFFSET_WIDTH, value); + } + + private void setElementSizeBuffer(int index, int value) { + sizeBuffer.setInt((long) index * SIZE_WIDTH, value); + } + private class TransferImpl implements TransferPair { ListViewVector to; @@ -499,7 +506,6 @@ public void splitAndTransfer(int startIndex, int length) { valueCount); to.clear(); if (length > 0) { - final int startPoint = offsetBuffer.getInt((long) startIndex * OFFSET_WIDTH); // we have to scan by index since there are out-of-order offsets to.offsetBuffer = to.allocateBuffers((long) length * OFFSET_WIDTH); to.sizeBuffer = to.allocateBuffers((long) length * SIZE_WIDTH); @@ -508,9 +514,9 @@ public void splitAndTransfer(int startIndex, int length) { int maxOffsetAndSizeSum = -1; int minOffsetValue = -1; for (int i = 0; i < length; i++) { - final int offsetValue = offsetBuffer.getInt((long) (startIndex + i) * OFFSET_WIDTH); - final int sizeValue = sizeBuffer.getInt((long) (startIndex + i) * SIZE_WIDTH); - to.sizeBuffer.setInt((long) i * SIZE_WIDTH, sizeValue); + final int offsetValue = getElementStartIndex(startIndex + i); + final int sizeValue = getElementSize(startIndex + i); + to.setElementSizeBuffer(i, sizeValue); if (maxOffsetAndSizeSum < offsetValue + sizeValue) { maxOffsetAndSizeSum = offsetValue + sizeValue; } @@ -521,9 +527,9 @@ public void splitAndTransfer(int startIndex, int length) { /* splitAndTransfer the offset buffer */ for (int i = 0; i < length; i++) { - final int offsetValue = offsetBuffer.getInt((long) (startIndex + i) * OFFSET_WIDTH); + final int offsetValue = getElementStartIndex(startIndex + i); final int relativeOffset = offsetValue - minOffsetValue; - to.offsetBuffer.setInt((long) i * OFFSET_WIDTH, relativeOffset); + to.setElementOffsetBuffer(i, relativeOffset); } /* splitAndTransfer the validity buffer */ @@ -536,70 +542,6 @@ public void splitAndTransfer(int startIndex, int length) { } } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer(int startIndex, int length, ListViewVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; @@ -634,7 +576,7 @@ public int getBufferSize() { } final int offsetBufferSize = valueCount * OFFSET_WIDTH; final int sizeBufferSize = valueCount * SIZE_WIDTH; - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return offsetBufferSize + sizeBufferSize + validityBufferSize + vector.getBufferSize(); } @@ -649,7 +591,7 @@ public int getBufferSizeFor(int valueCount) { if (valueCount == 0) { return 0; } - final int validityBufferSize = getValidityBufferSizeFromCount(valueCount); + final int validityBufferSize = BitVectorHelper.getValidityBufferSizeFromCount(valueCount); return super.getBufferSizeFor(valueCount) + validityBufferSize; } @@ -743,10 +685,10 @@ public List getObject(int index) { if (isSet(index) == 0) { return null; } - final List vals = new JsonStringArrayList<>(); - final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); - final int end = start + sizeBuffer.getInt((index) * SIZE_WIDTH); + final int start = getElementStartIndex(index); + final int end = getElementEndIndex(index); final ValueVector vv = getDataVector(); + final List vals = new JsonStringArrayList<>(end - start); for (int i = start; i < end; i++) { vals.add(vv.getObject(i)); } @@ -776,7 +718,7 @@ public boolean isEmpty(int index) { if (isNull(index)) { return true; } else { - return sizeBuffer.getInt(index * SIZE_WIDTH) == 0; + return getElementSize(index) == 0; } } @@ -787,10 +729,7 @@ public boolean isEmpty(int index) { * @return 1 if element at given index is not null, 0 otherwise */ public int isSet(int index) { - final int byteIndex = index >> 3; - final byte b = validityBuffer.getByte(byteIndex); - final int bitIndex = index & 7; - return (b >> bitIndex) & 0x01; + return BitVectorHelper.get(validityBuffer, index); } /** @@ -840,8 +779,8 @@ public void setNull(int index) { reallocValidityAndSizeAndOffsetBuffers(); } - offsetBuffer.setInt(index * OFFSET_WIDTH, 0); - sizeBuffer.setInt(index * SIZE_WIDTH, 0); + setElementOffsetBuffer(index, 0); + setElementSizeBuffer(index, 0); BitVectorHelper.unsetBit(validityBuffer, index); } @@ -859,11 +798,11 @@ public int startNewValue(int index) { if (index > 0) { final int prevOffset = getMaxViewEndChildVectorByIndex(index); - offsetBuffer.setInt(index * OFFSET_WIDTH, prevOffset); + setElementOffsetBuffer(index, prevOffset); } BitVectorHelper.setBit(validityBuffer, index); - return offsetBuffer.getInt(index * OFFSET_WIDTH); + return getElementStartIndex(index); } /** @@ -901,9 +840,9 @@ private void validateInvariants(int offset, int size) { * @param value value to set */ public void setOffset(int index, int value) { - validateInvariants(value, sizeBuffer.getInt(index * SIZE_WIDTH)); + validateInvariants(value, getElementSize(index)); - offsetBuffer.setInt(index * OFFSET_WIDTH, value); + setElementOffsetBuffer(index, value); } /** @@ -913,9 +852,9 @@ public void setOffset(int index, int value) { * @param value value to set */ public void setSize(int index, int value) { - validateInvariants(offsetBuffer.getInt(index * SIZE_WIDTH), value); + validateInvariants(getElementStartIndex(index), value); - sizeBuffer.setInt(index * SIZE_WIDTH, value); + setElementSizeBuffer(index, value); } /** @@ -951,12 +890,16 @@ public void setValueCount(int valueCount) { @Override public int getElementStartIndex(int index) { - return offsetBuffer.getInt(index * OFFSET_WIDTH); + return offsetBuffer.getInt((long) index * OFFSET_WIDTH); + } + + private int getElementSize(int index) { + return sizeBuffer.getInt((long) index * SIZE_WIDTH); } @Override public int getElementEndIndex(int index) { - return sizeBuffer.getInt(index * OFFSET_WIDTH); + return getElementStartIndex(index) + getElementSize(index); } @Override @@ -1013,8 +956,8 @@ public double getDensity() { @Override public void validate() { for (int i = 0; i < valueCount; i++) { - final int offset = offsetBuffer.getInt(i * OFFSET_WIDTH); - final int size = sizeBuffer.getInt(i * SIZE_WIDTH); + final int offset = getElementStartIndex(i); + final int size = getElementSize(i); validateInvariants(offset, size); } } @@ -1026,6 +969,6 @@ public void validate() { * @param size number of elements in the list that was written */ public void endValue(int index, int size) { - sizeBuffer.setInt(index * SIZE_WIDTH, size); + setElementSizeBuffer(index, size); } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java index 23cda8401b..3f98322ba9 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/MapVector.java @@ -22,7 +22,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.AddOrGetResult; -import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.ZeroVector; @@ -232,70 +231,6 @@ public void splitAndTransfer(int startIndex, int length) { } } - /* - * transfer the validity. - */ - private void splitAndTransferValidityBuffer(int startIndex, int length, MapVector target) { - int firstByteSource = BitVectorHelper.byteIndex(startIndex); - int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = getValidityBufferSizeFromCount(length); - int offset = startIndex % 8; - - if (length > 0) { - if (offset == 0) { - // slice - if (target.validityBuffer != null) { - target.validityBuffer.getReferenceManager().release(); - } - target.validityBuffer = validityBuffer.slice(firstByteSource, byteSizeTarget); - target.validityBuffer.getReferenceManager().retain(1); - } else { - /* Copy data - * When the first bit starts from the middle of a byte (offset != 0), - * copy data from src BitVector. - * Each byte in the target is composed by a part in i-th byte, - * another part in (i+1)-th byte. - */ - target.allocateValidityBuffer(byteSizeTarget); - - for (int i = 0; i < byteSizeTarget - 1; i++) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte(validityBuffer, firstByteSource + i, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + i + 1, offset); - - target.validityBuffer.setByte(i, (b1 + b2)); - } - - /* Copying the last piece is done in the following manner: - * if the source vector has 1 or more bytes remaining, we copy - * the last piece as a byte formed by shifting data - * from the current byte and the next byte. - * - * if the source vector has no more bytes remaining - * (we are at the last byte), we copy the last piece as a byte - * by shifting data from the current byte. - */ - if ((firstByteSource + byteSizeTarget - 1) < lastByteSource) { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - byte b2 = - BitVectorHelper.getBitsFromNextByte( - validityBuffer, firstByteSource + byteSizeTarget, offset); - - target.validityBuffer.setByte(byteSizeTarget - 1, b1 + b2); - } else { - byte b1 = - BitVectorHelper.getBitsFromCurrentByte( - validityBuffer, firstByteSource + byteSizeTarget - 1, offset); - target.validityBuffer.setByte(byteSizeTarget - 1, b1); - } - } - } - } - @Override public ValueVector getTo() { return to; diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java index 1bb9a3d6c0..b83e13449a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/RunEndEncodedVector.java @@ -28,6 +28,7 @@ import org.apache.arrow.memory.OutOfMemoryException; import org.apache.arrow.memory.util.ByteFunctionHelpers; import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BigIntVector; @@ -820,4 +821,101 @@ static int getPhysicalIndex(FieldVector runEndVector, int logicalIndex) { return result; } + + public static class RangeIterator { + + private final RunEndEncodedVector runEndEncodedVector; + private final int rangeEnd; + private int runIndex; + private int runEnd; + private int logicalPos; + + /** + * Constructs a new RangeIterator for iterating over a range of values in a RunEndEncodedVector. + * + * @param runEndEncodedVector The vector to iterate over + * @param startIndex The logical start index of the range (inclusive) + * @param length The number of values to include in the range + * @throws IllegalArgumentException if startIndex is negative or (startIndex + length) exceeds + * vector bounds + */ + public RangeIterator(RunEndEncodedVector runEndEncodedVector, int startIndex, int length) { + int rangeEnd = startIndex + length; + Preconditions.checkArgument( + startIndex >= 0, "startIndex %s must be non negative.", startIndex); + Preconditions.checkArgument( + rangeEnd <= runEndEncodedVector.getValueCount(), + "(startIndex + length) %s out of range[0, %s].", + rangeEnd, + runEndEncodedVector.getValueCount()); + + this.rangeEnd = rangeEnd; + this.runEndEncodedVector = runEndEncodedVector; + this.runIndex = runEndEncodedVector.getPhysicalIndex(startIndex) - 1; + this.runEnd = startIndex; + this.logicalPos = -1; + } + + /** + * Advances to the next run in the range. + * + * @return true if there is another run available, false if iteration has completed + */ + public boolean nextRun() { + logicalPos = runEnd; + if (logicalPos >= rangeEnd) { + return false; + } + updateRun(); + return true; + } + + private void updateRun() { + runIndex++; + runEnd = (int) ((BaseIntVector) runEndEncodedVector.runEndsVector).getValueAsLong(runIndex); + } + + /** + * Advances to the next value in the range. + * + * @return true if there is another value available, false if iteration has completed + */ + public boolean nextValue() { + logicalPos++; + if (logicalPos >= rangeEnd) { + return false; + } + if (logicalPos == runEnd) { + updateRun(); + } + return true; + } + + /** + * Gets the current run index (physical position in the run-ends vector). + * + * @return the current run index + */ + public int getRunIndex() { + return runIndex; + } + + /** + * Gets the length of the current run within the iterator's range. + * + * @return the number of remaining values in current run within the iterator's range + */ + public int getRunLength() { + return Math.min(runEnd, rangeEnd) - logicalPos; + } + + /** + * Checks if iteration has completed. + * + * @return true if all values in the range have been processed, false otherwise + */ + public boolean isEnd() { + return logicalPos >= rangeEnd; + } + } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java index ca5f572034..5e5bb7fc21 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/StructVector.java @@ -18,6 +18,7 @@ import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.apache.arrow.util.Preconditions.checkNotNull; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import java.util.ArrayList; import java.util.Arrays; @@ -89,7 +90,7 @@ public StructVector( super(name, checkNotNull(allocator), fieldType, callBack); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = - BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION); + getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** @@ -118,7 +119,7 @@ public StructVector( allowConflictPolicyChanges); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = - BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION); + getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** @@ -132,7 +133,7 @@ public StructVector(Field field, BufferAllocator allocator, CallBack callBack) { super(field, checkNotNull(allocator), callBack); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = - BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION); + getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } /** @@ -153,7 +154,7 @@ public StructVector( super(field, checkNotNull(allocator), callBack, conflictPolicy, allowConflictPolicyChanges); this.validityBuffer = allocator.getEmpty(); this.validityAllocationSizeInBytes = - BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION); + getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION); } @Override @@ -182,7 +183,7 @@ public List getFieldBuffers() { private void setReaderAndWriterIndex() { validityBuffer.readerIndex(0); - validityBuffer.writerIndex(BitVectorHelper.getValidityBufferSize(valueCount)); + validityBuffer.writerIndex(getValidityBufferSizeFromCount(valueCount)); } /** @@ -318,7 +319,7 @@ public void splitAndTransfer(int startIndex, int length) { private void splitAndTransferValidityBuffer(int startIndex, int length, StructVector target) { int firstByteSource = BitVectorHelper.byteIndex(startIndex); int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1); - int byteSizeTarget = BitVectorHelper.getValidityBufferSize(length); + int byteSizeTarget = getValidityBufferSizeFromCount(length); int offset = startIndex % 8; if (length > 0) { @@ -464,7 +465,7 @@ public int getBufferSize() { if (valueCount == 0) { return 0; } - return super.getBufferSize() + BitVectorHelper.getValidityBufferSize(valueCount); + return super.getBufferSize() + getValidityBufferSizeFromCount(valueCount); } /** @@ -478,18 +479,18 @@ public int getBufferSizeFor(final int valueCount) { if (valueCount == 0) { return 0; } - return super.getBufferSizeFor(valueCount) + BitVectorHelper.getValidityBufferSize(valueCount); + return super.getBufferSizeFor(valueCount) + getValidityBufferSizeFromCount(valueCount); } @Override public void setInitialCapacity(int numRecords) { - validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSize(numRecords); + validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords); } @Override public void setInitialCapacity(int numRecords, double density) { - validityAllocationSizeInBytes = BitVectorHelper.getValidityBufferSize(numRecords); + validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords); super.setInitialCapacity(numRecords, density); } @@ -547,7 +548,7 @@ private long getNewAllocationSize(int currentBufferCapacity) { newAllocationSize = validityAllocationSizeInBytes; } else { newAllocationSize = - BitVectorHelper.getValidityBufferSize(BaseValueVector.INITIAL_VALUE_ALLOCATION) * 2L; + getValidityBufferSizeFromCount(BaseValueVector.INITIAL_VALUE_ALLOCATION) * 2L; } } newAllocationSize = CommonUtil.nextPowerOfTwo(newAllocationSize); diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java deleted file mode 100644 index 09f0314c5f..0000000000 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/ExtensionTypeWriterFactory.java +++ /dev/null @@ -1,38 +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.arrow.vector.complex.impl; - -import org.apache.arrow.vector.ExtensionTypeVector; -import org.apache.arrow.vector.complex.writer.FieldWriter; - -/** - * A factory interface for creating instances of {@link ExtensionTypeWriter}. This factory allows - * configuring writer implementations for specific {@link ExtensionTypeVector}. - * - * @param the type of writer implementation for a specific {@link ExtensionTypeVector}. - */ -public interface ExtensionTypeWriterFactory { - - /** - * Returns an instance of the writer implementation for the given {@link ExtensionTypeVector}. - * - * @param vector the {@link ExtensionTypeVector} for which the writer implementation is to be - * returned. - * @return an instance of the writer implementation for the given {@link ExtensionTypeVector}. - */ - T getWriterImpl(ExtensionTypeVector vector); -} diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java new file mode 100644 index 0000000000..7a5312f6ed --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java @@ -0,0 +1,123 @@ +/* + * 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.arrow.vector.complex.impl; + +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.apache.arrow.vector.holders.UuidHolder; +import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.util.UuidUtility; + +/** + * Reader implementation for reading UUID values from a {@link NullableUuidHolder}. + * + *

This reader wraps a single UUID holder value and provides methods to read from it. Unlike + * {@link UuidReaderImpl} which reads from a vector, this reader operates on a holder instance. + * + * @see NullableUuidHolder + * @see UuidReaderImpl + */ +public class NullableUuidHolderReaderImpl extends AbstractFieldReader { + private final NullableUuidHolder holder; + + /** + * Constructs a reader for the given UUID holder. + * + * @param holder the UUID holder to read from + */ + public NullableUuidHolderReaderImpl(NullableUuidHolder holder) { + this.holder = holder; + } + + @Override + public int size() { + throw new UnsupportedOperationException( + "size() is not supported on NullableUuidHolderReaderImpl. " + + "This reader wraps a single UUID holder value, not a collection. " + + "Use UuidReaderImpl for vector-based UUID reading."); + } + + @Override + public boolean next() { + throw new UnsupportedOperationException( + "next() is not supported on NullableUuidHolderReaderImpl. " + + "This reader wraps a single UUID holder value, not an iterator. " + + "Use UuidReaderImpl for vector-based UUID reading."); + } + + @Override + public void setPosition(int index) { + throw new UnsupportedOperationException( + "setPosition() is not supported on NullableUuidHolderReaderImpl. " + + "This reader wraps a single UUID holder value, not a vector. " + + "Use UuidReaderImpl for vector-based UUID reading."); + } + + @Override + public Types.MinorType getMinorType() { + return Types.MinorType.EXTENSIONTYPE; + } + + @Override + public boolean isSet() { + return holder.isSet == 1; + } + + @Override + public void read(ExtensionHolder h) { + if (h instanceof NullableUuidHolder) { + NullableUuidHolder nullableHolder = (NullableUuidHolder) h; + nullableHolder.buffer = this.holder.buffer; + nullableHolder.isSet = this.holder.isSet; + nullableHolder.start = this.holder.start; + } else if (h instanceof UuidHolder) { + UuidHolder uuidHolder = (UuidHolder) h; + uuidHolder.buffer = this.holder.buffer; + uuidHolder.start = this.holder.start; + } else { + throw new IllegalArgumentException( + "Unsupported holder type: " + + h.getClass().getName() + + ". " + + "Only NullableUuidHolder and UuidHolder are supported for UUID values. " + + "Provided holder type cannot be used to read UUID data."); + } + } + + @Override + public Object readObject() { + if (!isSet()) { + return null; + } + // Convert UUID bytes to Java UUID object + try { + return UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + } catch (Exception e) { + throw new RuntimeException( + String.format( + "Failed to read UUID from buffer. Invalid Arrow buffer state: " + + "capacity=%d, readableBytes=%d, readerIndex=%d, writerIndex=%d, refCnt=%d. " + + "The buffer must contain exactly 16 bytes of valid UUID data.", + holder.buffer.capacity(), + holder.buffer.readableBytes(), + holder.buffer.readerIndex(), + holder.buffer.writerIndex(), + holder.buffer.refCnt()), + e); + } + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java index d341384bd9..93796aa77e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java @@ -60,11 +60,6 @@ public void writeExtension(Object var1) { } @Override - public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) { - this.writer = factory.getWriterImpl(vector); - this.writer.setPosition(idx()); - } - public void write(ExtensionHolder holder) { this.writer.write(holder); } @@ -76,4 +71,10 @@ public void setPosition(int index) { this.writer.setPosition(index); } } + + @Override + public void writeNull() { + this.vector.setNull(getPosition()); + this.vector.setValueCount(getPosition() + 1); + } } diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java new file mode 100644 index 0000000000..bb7ae13e5b --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java @@ -0,0 +1,94 @@ +/* + * 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.arrow.vector.complex.impl; + +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.apache.arrow.vector.holders.UuidHolder; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.Field; + +/** + * Reader implementation for {@link UuidVector}. + * + *

Provides methods to read UUID values from a vector, including support for reading into {@link + * UuidHolder} and retrieving values as {@link java.util.UUID} objects. + * + * @see UuidVector + * @see org.apache.arrow.vector.extension.UuidType + */ +public class UuidReaderImpl extends AbstractFieldReader { + + private final UuidVector vector; + + /** + * Constructs a reader for the given UUID vector. + * + * @param vector the UUID vector to read from + */ + public UuidReaderImpl(UuidVector vector) { + super(); + this.vector = vector; + } + + @Override + public MinorType getMinorType() { + return vector.getMinorType(); + } + + @Override + public Field getField() { + return vector.getField(); + } + + @Override + public boolean isSet() { + return !vector.isNull(idx()); + } + + @Override + public void read(ExtensionHolder holder) { + if (holder instanceof NullableUuidHolder) { + vector.get(idx(), (NullableUuidHolder) holder); + } else { + throw new IllegalArgumentException( + "Unsupported holder type for UuidReader: " + holder.getClass()); + } + } + + @Override + public void read(int arrayIndex, ExtensionHolder holder) { + if (holder instanceof NullableUuidHolder) { + vector.get(arrayIndex, (NullableUuidHolder) holder); + } else { + throw new IllegalArgumentException( + "Unsupported holder type for UuidReader: " + holder.getClass()); + } + } + + @Override + public void copyAsValue(AbstractExtensionTypeWriter writer) { + UuidWriterImpl impl = (UuidWriterImpl) writer; + impl.vector.copyFromSafe(idx(), impl.idx(), vector); + } + + @Override + public Object readObject() { + return vector.getObject(idx()); + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java new file mode 100644 index 0000000000..944b7e2e62 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.complex.impl; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.apache.arrow.vector.holders.UuidHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** + * Writer implementation for {@link UuidVector}. + * + *

Supports writing UUID values in multiple formats: {@link java.util.UUID}, byte arrays, and + * {@link ArrowBuf}. Also handles {@link UuidHolder} and {@link NullableUuidHolder}. + * + * @see UuidVector + * @see org.apache.arrow.vector.extension.UuidType + */ +public class UuidWriterImpl extends AbstractExtensionTypeWriter { + + /** + * Constructs a writer for the given UUID vector. + * + * @param vector the UUID vector to write to + */ + public UuidWriterImpl(UuidVector vector) { + super(vector); + } + + @Override + public void writeExtension(Object value) { + if (value instanceof byte[]) { + vector.setSafe(getPosition(), (byte[]) value); + } else if (value instanceof ArrowBuf) { + vector.setSafe(getPosition(), (ArrowBuf) value); + } else if (value instanceof java.util.UUID) { + vector.setSafe(getPosition(), (java.util.UUID) value); + } else if (value instanceof ExtensionHolder) { + write((ExtensionHolder) value); + } else { + throw new IllegalArgumentException( + "Unsupported value type for UUID: " + + value.getClass().getName() + + ". " + + "Supported types are: byte[] (16 bytes), ArrowBuf (16 bytes), or java.util.UUID. " + + "Convert your value to one of these types before writing."); + } + vector.setValueCount(getPosition() + 1); + } + + @Override + public void writeExtension(Object value, ArrowType type) { + writeExtension(value); + } + + @Override + public void write(ExtensionHolder holder) { + if (holder instanceof UuidHolder) { + vector.setSafe(getPosition(), (UuidHolder) holder); + } else if (holder instanceof NullableUuidHolder) { + vector.setSafe(getPosition(), (NullableUuidHolder) holder); + } + vector.setValueCount(getPosition() + 1); + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java b/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java new file mode 100644 index 0000000000..1ba7b27156 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/reader/ExtensionReader.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.complex.reader; + +import org.apache.arrow.vector.holders.ExtensionHolder; + +/** Interface for reading extension types. Extends the functionality of {@link BaseReader}. */ +public interface ExtensionReader extends BaseReader { + + /** + * Reads to the given extension holder. + * + * @param holder the {@link ExtensionHolder} to read + */ + void read(ExtensionHolder holder); + + /** + * Reads and returns an object representation of the extension type. + * + * @return the object representation of the extension type + */ + Object readObject(); + + /** + * Checks if the current value is set. + * + * @return true if the value is set, false otherwise + */ + boolean isSet(); +} diff --git a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java b/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java index 58d9e4db9b..b108173c82 100644 --- a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java +++ b/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java @@ -29,7 +29,11 @@ public abstract class AbstractCompressionCodec implements CompressionCodec { @Override public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) { - if (uncompressedBuffer.writerIndex() == 0L) { + // GH-1116: capture writerIndex() once so the empty-buffer check, size + // comparison, and uncompressed-length prefix all see the same value. + long uncompressedLength = uncompressedBuffer.writerIndex(); + + if (uncompressedLength == 0L) { // shortcut for empty buffer ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH); compressedBuffer.setLong(0, 0); @@ -41,7 +45,6 @@ public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer); long compressedLength = compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH; - long uncompressedLength = uncompressedBuffer.writerIndex(); if (compressedLength > uncompressedLength) { // compressed buffer is larger, send the raw buffer diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java index ca56214fda..780a4ee659 100644 --- a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java +++ b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java @@ -54,10 +54,12 @@ import org.apache.arrow.vector.TimeStampNanoVector; import org.apache.arrow.vector.TimeStampSecTZVector; import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.ViewVarBinaryVector; import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; @@ -177,6 +179,11 @@ public int hashCode() { return Objects.hash(super.hashCode(), storageType, typeName, vendorName); } + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + throw new UnsupportedOperationException("WriterImpl not yet implemented."); + } + @Override public String toString() { return "OpaqueType(" diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java b/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java new file mode 100644 index 0000000000..c249c6eda9 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java @@ -0,0 +1,119 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.extension; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.complex.impl.UuidWriterImpl; +import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; +import org.apache.arrow.vector.types.pojo.FieldType; + +/** + * Extension type for UUID (Universally Unique Identifier) values. + * + *

UUIDs are stored as 16-byte fixed-size binary values. This extension type provides a + * standardized way to represent UUIDs in Arrow, making them interoperable across different systems + * and languages.Ï€ + * + *

The extension name is "arrow.uuid" and it uses {@link ArrowType.FixedSizeBinary} with 16 bytes + * as the storage type. + * + *

Usage: + * + * 

{@code
+ * UuidVector vector = new UuidVector("uuid_col", allocator);
+ * vector.set(0, UUID.randomUUID());
+ * UUID value = vector.getObject(0);
+ * }
+ * + * @see UuidVector + * @see org.apache.arrow.vector.holders.UuidHolder + * @see org.apache.arrow.vector.holders.NullableUuidHolder + */ +public class UuidType extends ExtensionType { + /** Singleton instance of UuidType. */ + public static final UuidType INSTANCE = new UuidType(); + + /** Extension name registered in the Arrow extension type registry. */ + public static final String EXTENSION_NAME = "arrow.uuid"; + + /** Number of bytes used to store a UUID (128 bits = 16 bytes). */ + public static final int UUID_BYTE_WIDTH = 16; + + /** Number of characters in the standard UUID string representation (with hyphens). */ + public static final int UUID_STRING_WIDTH = 36; + + /** Storage type for UUID: FixedSizeBinary(16). */ + public static final ArrowType STORAGE_TYPE = new ArrowType.FixedSizeBinary(UUID_BYTE_WIDTH); + + private UuidType() {} + + static { + ExtensionTypeRegistry.register(INSTANCE); + } + + @Override + public ArrowType storageType() { + return STORAGE_TYPE; + } + + @Override + public String extensionName() { + return EXTENSION_NAME; + } + + @Override + public boolean extensionEquals(ExtensionType other) { + return other instanceof UuidType; + } + + @Override + public ArrowType deserialize(ArrowType storageType, String serializedData) { + if (!storageType.equals(storageType())) { + throw new UnsupportedOperationException( + "Cannot construct UuidType from underlying type " + storageType); + } + return INSTANCE; + } + + @Override + public String serialize() { + return ""; + } + + @Override + public boolean isComplex() { + return false; + } + + @Override + public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { + return new UuidVector( + name, fieldType, allocator, new FixedSizeBinaryVector(name, allocator, UUID_BYTE_WIDTH)); + } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new UuidWriterImpl((UuidVector) vector); + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java index fc7ed85878..4d3f767aef 100644 --- a/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java +++ b/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java @@ -16,7 +16,11 @@ */ package org.apache.arrow.vector.holders; +import org.apache.arrow.vector.types.pojo.ArrowType; + /** Base {@link ValueHolder} class for a {@link org.apache.arrow.vector.ExtensionTypeVector}. */ public abstract class ExtensionHolder implements ValueHolder { public int isSet; + + public abstract ArrowType type(); } diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java new file mode 100644 index 0000000000..6a2b4ff604 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.holders; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** + * Value holder for nullable UUID values. + * + *

The {@code isSet} field controls nullability: when {@code isSet = 1}, the holder contains a + * valid UUID in {@code buffer}; when {@code isSet = 0}, the holder represents a null value and + * {@code buffer} should not be accessed. + * + * @see UuidHolder + * @see org.apache.arrow.vector.UuidVector + * @see org.apache.arrow.vector.extension.UuidType + */ +public class NullableUuidHolder extends ExtensionHolder { + /** Buffer containing 16-byte UUID data. */ + public ArrowBuf buffer; + + /** Offset in the buffer where the UUID data starts. */ + public int start = 0; + + @Override + public ArrowType type() { + return UuidType.INSTANCE; + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java b/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java new file mode 100644 index 0000000000..9ec0305f30 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java @@ -0,0 +1,48 @@ +/* + * 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.arrow.vector.holders; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** + * Value holder for non-nullable UUID values. + * + *

Contains a 16-byte UUID in {@code buffer} with {@code isSet} always 1. + * + * @see NullableUuidHolder + * @see org.apache.arrow.vector.UuidVector + * @see org.apache.arrow.vector.extension.UuidType + */ +public class UuidHolder extends ExtensionHolder { + /** Buffer containing 16-byte UUID data. */ + public ArrowBuf buffer; + + /** Offset in the buffer where the UUID data starts. */ + public int start = 0; + + /** Constructs a UuidHolder with isSet = 1. */ + public UuidHolder() { + this.isSet = 1; + } + + @Override + public ArrowType type() { + return UuidType.INSTANCE; + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java index fe0803d298..e4bab7eb80 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java @@ -20,6 +20,7 @@ import static com.fasterxml.jackson.core.JsonToken.END_OBJECT; import static com.fasterxml.jackson.core.JsonToken.START_ARRAY; import static com.fasterxml.jackson.core.JsonToken.START_OBJECT; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.apache.arrow.vector.BufferLayout.BufferType.DATA; import static org.apache.arrow.vector.BufferLayout.BufferType.OFFSET; import static org.apache.arrow.vector.BufferLayout.BufferType.SIZE; @@ -381,7 +382,7 @@ private class BufferHelper { new BufferReader() { @Override protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { - final int bufferSize = BitVectorHelper.getValidityBufferSize(count); + final int bufferSize = getValidityBufferSizeFromCount(count); ArrowBuf buf = allocator.buffer(bufferSize); // C++ integration test fails without this. diff --git a/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java b/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java new file mode 100644 index 0000000000..a1b0b54579 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/util/UuidUtility.java @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector.util; + +import static org.apache.arrow.vector.extension.UuidType.UUID_BYTE_WIDTH; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.UUID; +import org.apache.arrow.memory.ArrowBuf; + +/** + * Utility class for UUID conversions and operations. + * + *

Provides methods to convert between {@link UUID} objects and byte representations used in + * Arrow vectors. + * + * @see org.apache.arrow.vector.UuidVector + * @see org.apache.arrow.vector.extension.UuidType + */ +public class UuidUtility { + /** + * Converts a UUID to a 16-byte array. + * + *

The UUID is stored in big-endian byte order, with the most significant bits first. + * + * @param uuid the UUID to convert + * @return a 16-byte array representing the UUID + */ + public static byte[] getBytesFromUUID(UUID uuid) { + byte[] result = new byte[16]; + long msb = uuid.getMostSignificantBits(); + long lsb = uuid.getLeastSignificantBits(); + for (int i = 15; i >= 8; i--) { + result[i] = (byte) (lsb & 0xFF); + lsb >>= 8; + } + for (int i = 7; i >= 0; i--) { + result[i] = (byte) (msb & 0xFF); + msb >>= 8; + } + return result; + } + + /** + * Constructs a UUID from bytes stored in an ArrowBuf at the specified index. + * + *

Reads 16 bytes from the buffer starting at the given index and interprets them as a UUID in + * big-endian byte order. + * + * @param buffer the buffer containing UUID data + * @param index the byte offset in the buffer where the UUID starts + * @return the UUID constructed from the buffer data + */ + public static UUID uuidFromArrowBuf(ArrowBuf buffer, long index) { + ByteBuffer buf = buffer.nioBuffer(index, UUID_BYTE_WIDTH); + + buf.order(ByteOrder.BIG_ENDIAN); + long mostSigBits = buf.getLong(0); + long leastSigBits = buf.getLong(Long.BYTES); + return new UUID(mostSigBits, leastSigBits); + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java index 0dc96a4d4b..2cfeb0a04d 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/VectorAppender.java @@ -24,13 +24,17 @@ import org.apache.arrow.memory.util.MemoryUtil; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseFixedWidthVector; +import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; +import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.ExtensionTypeVector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.NullVector; +import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.compare.TypeEqualsVisitor; import org.apache.arrow.vector.compare.VectorVisitor; @@ -39,6 +43,7 @@ import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.NonNullableStructVector; +import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.UnionVector; /** Utility to append two vectors together. */ @@ -120,10 +125,15 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { targetVector .getOffsetBuffer() .getInt((long) targetVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); + // The delta vector's offset buffer need not start at zero (e.g. a vector imported through + // the C data interface from a sliced array), so the amount of data to append is the + // distance between its first and last offsets, not the last offset itself. + int deltaDataStart = deltaVector.getOffsetBuffer().getInt(0); int deltaDataSize = deltaVector - .getOffsetBuffer() - .getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); + .getOffsetBuffer() + .getInt((long) deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH) + - deltaDataStart; int newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity @@ -144,7 +154,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { // append data buffer MemoryUtil.copyMemory( - deltaVector.getDataBuffer().memoryAddress(), + deltaVector.getDataBuffer().memoryAddress() + deltaDataStart, targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); @@ -155,7 +165,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { + (targetVector.getValueCount() + 1) * BaseVariableWidthVector.OFFSET_WIDTH, deltaVector.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH); - // increase each offset from the second buffer + // rebase each appended offset to the target's data, accounting for the delta's start offset for (int i = 0; i < deltaVector.getValueCount(); i++) { int oldOffset = targetVector @@ -167,7 +177,7 @@ public ValueVector visit(BaseVariableWidthVector deltaVector, Void value) { .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * BaseVariableWidthVector.OFFSET_WIDTH, - oldOffset + targetDataSize); + oldOffset - deltaDataStart + targetDataSize); } ((BaseVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); @@ -191,11 +201,15 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { .getOffsetBuffer() .getLong( (long) targetVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); + // see the corresponding comment in visit(BaseVariableWidthVector, Void): the delta's + // offset buffer need not start at zero + long deltaDataStart = deltaVector.getOffsetBuffer().getLong(0); long deltaDataSize = deltaVector - .getOffsetBuffer() - .getLong( - (long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); + .getOffsetBuffer() + .getLong( + (long) deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH) + - deltaDataStart; long newValueCapacity = targetDataSize + deltaDataSize; // make sure there is enough capacity @@ -216,7 +230,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { // append data buffer MemoryUtil.copyMemory( - deltaVector.getDataBuffer().memoryAddress(), + deltaVector.getDataBuffer().memoryAddress() + deltaDataStart, targetVector.getDataBuffer().memoryAddress() + targetDataSize, deltaDataSize); @@ -227,7 +241,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { + (targetVector.getValueCount() + 1) * BaseLargeVariableWidthVector.OFFSET_WIDTH, deltaVector.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH); - // increase each offset from the second buffer + // rebase each appended offset to the target's data, accounting for the delta's start offset for (int i = 0; i < deltaVector.getValueCount(); i++) { long oldOffset = targetVector @@ -240,7 +254,7 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { .setLong( (long) (targetVector.getValueCount() + 1 + i) * BaseLargeVariableWidthVector.OFFSET_WIDTH, - oldOffset + targetDataSize); + oldOffset - deltaDataStart + targetDataSize); } ((BaseLargeVariableWidthVector) targetVector).setLastSet(newValueCount - 1); targetVector.setValueCount(newValueCount); @@ -326,16 +340,20 @@ public ValueVector visit(ListVector deltaVector, Void value) { targetVector .getOffsetBuffer() .getInt((long) targetVector.getValueCount() * ListVector.OFFSET_WIDTH); - int deltaListSize = + // see the corresponding comment in visit(BaseVariableWidthVector, Void): the delta's + // offset buffer need not start at zero + int deltaListStart = deltaVector.getOffsetBuffer().getInt(0); + int deltaListEnd = deltaVector .getOffsetBuffer() .getInt((long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); + int deltaListSize = deltaListEnd - deltaListStart; ListVector targetListVector = (ListVector) targetVector; // make sure the underlying vector has value count set targetListVector.getDataVector().setValueCount(targetListSize); - deltaVector.getDataVector().setValueCount(deltaListSize); + deltaVector.getDataVector().setValueCount(deltaListEnd); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { @@ -367,13 +385,16 @@ public ValueVector visit(ListVector deltaVector, Void value) { .getOffsetBuffer() .setInt( (long) (targetVector.getValueCount() + 1 + i) * ListVector.OFFSET_WIDTH, - oldOffset + targetListSize); + oldOffset - deltaListStart + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors - VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); - deltaVector.getDataVector().accept(innerAppender, null); + appendDataVector( + targetListVector.getDataVector(), + deltaVector.getDataVector(), + deltaListStart, + deltaListSize); targetVector.setValueCount(newValueCount); return targetVector; @@ -395,17 +416,21 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { targetVector .getOffsetBuffer() .getLong((long) targetVector.getValueCount() * LargeListVector.OFFSET_WIDTH); - long deltaListSize = + // see the corresponding comment in visit(BaseVariableWidthVector, Void): the delta's + // offset buffer need not start at zero + long deltaListStart = deltaVector.getOffsetBuffer().getLong(0); + long deltaListEnd = deltaVector .getOffsetBuffer() .getLong((long) deltaVector.getValueCount() * LargeListVector.OFFSET_WIDTH); + long deltaListSize = deltaListEnd - deltaListStart; - ListVector targetListVector = (ListVector) targetVector; + LargeListVector targetListVector = (LargeListVector) targetVector; // make sure the underlying vector has value count set // todo recheck these casts when int64 vectors are supported targetListVector.getDataVector().setValueCount(checkedCastToInt(targetListSize)); - deltaVector.getDataVector().setValueCount(checkedCastToInt(deltaListSize)); + deltaVector.getDataVector().setValueCount(checkedCastToInt(deltaListEnd)); // make sure there is enough capacity while (targetVector.getValueCapacity() < newValueCount) { @@ -422,10 +447,10 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { // append offset buffer MemoryUtil.copyMemory( - deltaVector.getOffsetBuffer().memoryAddress() + ListVector.OFFSET_WIDTH, + deltaVector.getOffsetBuffer().memoryAddress() + LargeListVector.OFFSET_WIDTH, targetVector.getOffsetBuffer().memoryAddress() + (targetVector.getValueCount() + 1) * LargeListVector.OFFSET_WIDTH, - (long) deltaVector.getValueCount() * ListVector.OFFSET_WIDTH); + (long) deltaVector.getValueCount() * LargeListVector.OFFSET_WIDTH); // increase each offset from the second buffer for (int i = 0; i < deltaVector.getValueCount(); i++) { @@ -438,18 +463,42 @@ public ValueVector visit(LargeListVector deltaVector, Void value) { .getOffsetBuffer() .setLong( (long) (targetVector.getValueCount() + 1 + i) * LargeListVector.OFFSET_WIDTH, - oldOffset + targetListSize); + oldOffset - deltaListStart + targetListSize); } targetListVector.setLastSet(newValueCount - 1); // append underlying vectors - VectorAppender innerAppender = new VectorAppender(targetListVector.getDataVector()); - deltaVector.getDataVector().accept(innerAppender, null); + appendDataVector( + targetListVector.getDataVector(), + deltaVector.getDataVector(), + checkedCastToInt(deltaListStart), + checkedCastToInt(deltaListSize)); targetVector.setValueCount(newValueCount); return targetVector; } + /** + * Appends the range [start, start + length) of the delta vector's data vector to the target + * vector's data vector. The range may not cover the whole delta data vector when the delta's + * offset buffer does not start at zero. + */ + private static void appendDataVector( + ValueVector targetDataVector, ValueVector deltaDataVector, int start, int length) { + if (start == 0 && length == deltaDataVector.getValueCount()) { + VectorAppender innerAppender = new VectorAppender(targetDataVector); + deltaDataVector.accept(innerAppender, null); + return; + } + TransferPair transferPair = + deltaDataVector.getTransferPair(deltaDataVector.getField(), deltaDataVector.getAllocator()); + transferPair.splitAndTransfer(start, length); + try (ValueVector slicedDeltaDataVector = transferPair.getTo()) { + VectorAppender innerAppender = new VectorAppender(targetDataVector); + slicedDeltaDataVector.accept(innerAppender, null); + } + } + @Override public ValueVector visit(FixedSizeListVector deltaVector, Void value) { Preconditions.checkArgument( @@ -698,4 +747,98 @@ public ValueVector visit(ExtensionTypeVector deltaVector, Void value) { deltaVector.getUnderlyingVector().accept(underlyingAppender, null); return targetVector; } + + @Override + public ValueVector visit(RunEndEncodedVector deltaVector, Void value) { + Preconditions.checkArgument( + typeVisitor.equals(deltaVector), + "The deltaVector to append must have the same type as the targetVector"); + + if (deltaVector.getValueCount() == 0) { + return targetVector; // optimization, nothing to append, return + } + + RunEndEncodedVector targetEncodedVector = (RunEndEncodedVector) targetVector; + + final int targetLogicalValueCount = targetEncodedVector.getValueCount(); + + // Append the values vector first. + VectorAppender valueAppender = new VectorAppender(targetEncodedVector.getValuesVector()); + deltaVector.getValuesVector().accept(valueAppender, null); + + // Then append the run-ends vector. + BaseIntVector targetRunEndsVector = (BaseIntVector) targetEncodedVector.getRunEndsVector(); + BaseIntVector deltaRunEndsVector = (BaseIntVector) deltaVector.getRunEndsVector(); + appendRunEndsVector(targetRunEndsVector, deltaRunEndsVector, targetLogicalValueCount); + + targetEncodedVector.setValueCount(targetLogicalValueCount + deltaVector.getValueCount()); + return targetVector; + } + + private void appendRunEndsVector( + BaseIntVector targetRunEndsVector, + BaseIntVector deltaRunEndsVector, + int targetLogicalValueCount) { + int targetPhysicalValueCount = targetRunEndsVector.getValueCount(); + int newPhysicalValueCount = targetPhysicalValueCount + deltaRunEndsVector.getValueCount(); + + // make sure there is enough capacity + while (targetVector.getValueCapacity() < newPhysicalValueCount) { + targetVector.reAlloc(); + } + + // append validity buffer + BitVectorHelper.concatBits( + targetRunEndsVector.getValidityBuffer(), + targetRunEndsVector.getValueCount(), + deltaRunEndsVector.getValidityBuffer(), + deltaRunEndsVector.getValueCount(), + targetRunEndsVector.getValidityBuffer()); + + // shift and append data buffer + shiftAndAppendRunEndsDataBuffer( + targetRunEndsVector, + targetPhysicalValueCount, + deltaRunEndsVector.getDataBuffer(), + targetLogicalValueCount, + deltaRunEndsVector.getValueCount()); + + targetRunEndsVector.setValueCount(newPhysicalValueCount); + } + + private void shiftAndAppendRunEndsDataBuffer( + BaseIntVector toRunEndVector, + int toIndex, + ArrowBuf fromRunEndBuffer, + int offset, + int physicalLength) { + ArrowBuf toRunEndBuffer = toRunEndVector.getDataBuffer(); + if (toRunEndVector instanceof SmallIntVector) { + byte typeWidth = SmallIntVector.TYPE_WIDTH; + for (int i = 0; i < physicalLength; i++) { + toRunEndBuffer.setShort( + (long) (i + toIndex) * typeWidth, + fromRunEndBuffer.getShort((long) (i) * typeWidth) + offset); + } + + } else if (toRunEndVector instanceof IntVector) { + byte typeWidth = IntVector.TYPE_WIDTH; + for (int i = 0; i < physicalLength; i++) { + toRunEndBuffer.setInt( + (long) (i + toIndex) * typeWidth, + fromRunEndBuffer.getInt((long) (i) * typeWidth) + offset); + } + + } else if (toRunEndVector instanceof BigIntVector) { + byte typeWidth = BigIntVector.TYPE_WIDTH; + for (int i = 0; i < physicalLength; i++) { + toRunEndBuffer.setLong( + (long) (i + toIndex) * typeWidth, + fromRunEndBuffer.getLong((long) (i) * typeWidth) + offset); + } + } else { + throw new IllegalArgumentException( + "Run-end vector and must be of type int with size 16, 32, or 64 bits."); + } + } } diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java index 5c7215437f..5cfe64b14e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorBufferVisitor.java @@ -52,14 +52,22 @@ private void validateVectorCommon(ValueVector vector) { if (vector instanceof FieldVector) { FieldVector fieldVector = (FieldVector) vector; - // TODO: https://github.com/apache/arrow/issues/41734 int typeBufferCount = TypeLayout.getTypeBufferCount(arrowType); - validateOrThrow( - fieldVector.getFieldBuffers().size() == typeBufferCount, - "Expected %s buffers in vector of type %s, got %s.", - typeBufferCount, - vector.getField().getType().toString(), - fieldVector.getFieldBuffers().size()); + if (TypeLayout.getTypeLayout(arrowType).isFixedBufferCount()) { + validateOrThrow( + fieldVector.getFieldBuffers().size() == typeBufferCount, + "Expected %s buffers in vector of type %s, got %s.", + typeBufferCount, + vector.getField().getType().toString(), + fieldVector.getFieldBuffers().size()); + } else { + validateOrThrow( + fieldVector.getFieldBuffers().size() >= typeBufferCount, + "Expected at least %s buffers in vector of type %s, got %s.", + typeBufferCount, + vector.getField().getType().toString(), + fieldVector.getFieldBuffers().size()); + } } } @@ -158,7 +166,12 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) { @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { - throw new UnsupportedOperationException("View vectors are not supported."); + final int valueCount = vector.getValueCount(); + validateVectorCommon(vector); + validateOrThrow(vector.getFieldBuffers().size() >= 2, "Expected at least 2 buffers."); + validateValidityBuffer(vector, valueCount); + validateDataBuffer(vector, (long) valueCount * BaseVariableWidthViewVector.ELEMENT_SIZE); + return null; } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java index c62bff79f7..9da8cc813e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorDataVisitor.java @@ -121,7 +121,8 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) { @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { - throw new UnsupportedOperationException("View vectors are not supported."); + vector.validateScalars(); + return null; } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java index daad41dbdc..395852ef79 100644 --- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java @@ -61,6 +61,8 @@ import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarBinaryVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.compare.VectorVisitor; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; @@ -380,7 +382,12 @@ public Void visit(BaseLargeVariableWidthVector vector, Void value) { @Override public Void visit(BaseVariableWidthViewVector vector, Void value) { - throw new UnsupportedOperationException("View vectors are not supported."); + if (vector instanceof ViewVarCharVector) { + validateVectorCommon(vector, ArrowType.Utf8View.class); + } else if (vector instanceof ViewVarBinaryVector) { + validateVectorCommon(vector, ArrowType.BinaryView.class); + } + return null; } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java index 5004ba488c..2111410016 100644 --- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorVisitor.java @@ -107,8 +107,13 @@ public Void visit(BaseLargeVariableWidthVector left, Void value) { } @Override - public Void visit(BaseVariableWidthViewVector left, Void value) { - throw new UnsupportedOperationException("View vectors are not supported."); + public Void visit(BaseVariableWidthViewVector vector, Void value) { + if (vector.getValueCount() > 0) { + if (vector.getDataBuffer() == null || vector.getDataBuffer().capacity() == 0) { + throw new IllegalArgumentException("valueBuffer is null or capacity is 0"); + } + } + return null; } @Override diff --git a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java index 73a88b3a1e..b3455fe52c 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java @@ -30,14 +30,21 @@ import java.util.Arrays; import java.util.List; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader; import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter; import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.holders.DurationHolder; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; +import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; +import org.apache.arrow.vector.holders.TimeStampNanoTZHolder; +import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.TransferPair; @@ -628,6 +635,206 @@ public void testWriteLargeVarBinaryHelpers() throws Exception { } } + @Test + public void testWriterTimeStampNanoTZField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + final int valueCount = 10; + + for (int i = 0; i < valueCount; i++) { + writer.startList(); + writer.timeStampNanoTZ().writeTimeStampNanoTZ(i * 1000L); + writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 1) * 1000L); + writer.timeStampNanoTZ().writeTimeStampNanoTZ((i + 2) * 1000L); + writer.endList(); + } + vector.setValueCount(valueCount); + + UnionFixedSizeListReader reader = vector.getReader(); + for (int i = 0; i < valueCount; i++) { + reader.setPosition(i); + assertTrue(reader.isSet()); + assertTrue(reader.next()); + assertEquals(i * 1000L, reader.reader().readLong().longValue()); + assertTrue(reader.next()); + assertEquals((i + 1) * 1000L, reader.reader().readLong().longValue()); + assertTrue(reader.next()); + assertEquals((i + 2) * 1000L, reader.reader().readLong().longValue()); + assertFalse(reader.next()); + } + } + } + + @Test + public void testWriterUsingHolderTimeStampNanoTZField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + TimeStampNanoTZHolder holder = new TimeStampNanoTZHolder(); + holder.timezone = "SomeFakeTimeZone"; + writer.startList(); + holder.value = 12341234L; + writer.timeStampNanoTZ().write(holder); + holder.value = 55555L; + writer.timeStampNanoTZ().write(holder); + + // Writing with a different timezone should throw + holder.timezone = "AsdfTimeZone"; + holder.value = 77777; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, () -> writer.timeStampNanoTZ().write(holder)); + assertEquals( + "holder.timezone: AsdfTimeZone not equal to vector timezone: SomeFakeTimeZone", + ex.getMessage()); + + writer.endList(); + vector.setValueCount(1); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "SomeFakeTimeZone")), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(3)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + + @Test + public void testWriterUsingHolderTimestampMilliTZField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + TimeStampMilliTZHolder holder = new TimeStampMilliTZHolder(); + holder.timezone = "SomeFakeTimeZone"; + writer.startList(); + holder.value = 12341234L; + writer.timeStampMilliTZ().write(holder); + holder.value = 55555L; + writer.timeStampMilliTZ().write(holder); + + // Writing with a different timezone should throw + holder.timezone = "AsdfTimeZone"; + holder.value = 77777; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, () -> writer.timeStampMilliTZ().write(holder)); + assertEquals( + "holder.timezone: AsdfTimeZone not equal to vector timezone: SomeFakeTimeZone", + ex.getMessage()); + + writer.endList(); + vector.setValueCount(1); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "SomeFakeTimeZone")), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(3)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + + @Test + public void testWriterUsingHolderDurationField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + DurationHolder durationHolder = new DurationHolder(); + durationHolder.unit = TimeUnit.MILLISECOND; + + writer.startList(); + durationHolder.value = 812374L; + writer.duration().write(durationHolder); + durationHolder.value = 143451L; + writer.duration().write(durationHolder); + + // Writing with a different unit should throw + durationHolder.unit = TimeUnit.SECOND; + durationHolder.value = 8888888; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, () -> writer.duration().write(durationHolder)); + assertEquals("holder.unit: SECOND not equal to vector unit: MILLISECOND", ex.getMessage()); + + writer.endList(); + vector.setValueCount(1); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.Duration(TimeUnit.MILLISECOND)), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(3)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + + @Test + public void testWriterUsingHolderFixedSizeBinaryField() { + try (final FixedSizeListVector vector = + FixedSizeListVector.empty("vector", /* size= */ 2, allocator)) { + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + + writer.startList(); + writer.fixedSizeBinary().write(holder1); + holder1.buffer.close(); + writer.fixedSizeBinary().write(holder2); + holder2.buffer.close(); + + writer.endList(); + vector.setValueCount(1); + + FieldReader reader = vector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + + Field expectedDataField = + new Field( + BaseRepeatedValueVector.DATA_VECTOR_NAME, + FieldType.nullable(new ArrowType.FixedSizeBinary(2)), + null); + Field expectedField = + new Field( + vector.getName(), + FieldType.nullable(new ArrowType.FixedSizeList(2)), + List.of(expectedDataField)); + + assertEquals(expectedField, writer.getField()); + } + } + private int[] convertListToIntArray(List list) { int[] values = new int[list.size()]; for (int i = 0; i < list.size(); i++) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java index 101d942d2a..bf9bba9c78 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; @@ -25,18 +26,24 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionLargeListReader; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.holders.NullableUuidHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -943,7 +950,7 @@ public void testGetBufferSizeFor() { int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { - int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount); + int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = (valueCount + 1) * LargeListVector.OFFSET_WIDTH; int expectedSize = @@ -1020,6 +1027,99 @@ public void testGetTransferPairWithField() throws Exception { } } + @Test + public void testCopyValueSafeForExtensionType() throws Exception { + try (LargeListVector inVector = LargeListVector.empty("input", allocator); + LargeListVector outVector = LargeListVector.empty("output", allocator)) { + UnionLargeListWriter writer = inVector.getWriter(); + writer.allocate(); + + // Create first list with UUIDs + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u2); + writer.endList(); + + // Create second list with UUIDs + writer.setPosition(1); + UUID u3 = UUID.randomUUID(); + UUID u4 = UUID.randomUUID(); + writer.startList(); + extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u3); + extensionWriter.writeExtension(u4); + extensionWriter.writeNull(); + + writer.endList(); + writer.setValueCount(2); + + // Use copyFromSafe with ExtensionTypeWriterFactory + // This internally calls TransferImpl.copyValueSafe with ExtensionTypeWriterFactory + outVector.allocateNew(); + TransferPair tp = inVector.makeTransferPair(outVector); + tp.copyValueSafe(0, 0); + tp.copyValueSafe(1, 1); + outVector.setValueCount(2); + + // Verify first list + UnionLargeListReader reader = outVector.getReader(); + reader.setPosition(0); + assertTrue(reader.isSet(), "first list shouldn't be null"); + reader.next(); + FieldReader uuidReader = reader.reader(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + + // Verify second list + reader.setPosition(1); + assertTrue(reader.isSet(), "second list shouldn't be null"); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u4, actualUuid); + reader.next(); + uuidReader = reader.reader(); + assertFalse(uuidReader.isSet(), "third element should be null"); + } + } + + @Test + public void testEmptyLargeListOffsetBuffer() { + // Test that LargeListVector has correct readableBytes after allocation. + // According to Arrow spec, offset buffer must have N+1 entries. + // Even when N=0, it should contain [0]. + try (LargeListVector list = LargeListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.allocateNew(); + list.setValueCount(0); + + List buffers = list.getFieldBuffers(); + assertTrue( + buffers.get(1).readableBytes() >= LargeListVector.OFFSET_WIDTH, + "Offset buffer should have at least " + + LargeListVector.OFFSET_WIDTH + + " bytes for offset[0]"); + assertEquals(0L, list.getOffsetBuffer().getLong(0)); + } + } + private void writeIntValues(UnionLargeListWriter writer, int[] values) { writer.startList(); for (int v : values) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java index 26e7bb4a0d..256aa99687 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestLargeListViewVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertSame; @@ -1062,7 +1063,7 @@ public void testGetBufferSizeFor() { int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { - int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount); + int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = valueCount * BaseLargeRepeatedValueViewVector.OFFSET_WIDTH; int sizeBufferSize = valueCount * BaseLargeRepeatedValueViewVector.SIZE_WIDTH; diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java index 1d6fa39f9e..0c90b32abc 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; @@ -26,15 +27,20 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListReader; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types.MinorType; @@ -42,6 +48,7 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -1123,7 +1130,7 @@ public void testGetBufferSizeFor() { int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { - int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount); + int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = (valueCount + 1) * BaseRepeatedValueVector.OFFSET_WIDTH; int expectedSize = @@ -1198,6 +1205,200 @@ public void testGetTransferPairWithField() { } } + @Test + public void testListVectorWithExtensionType() throws Exception { + final FieldType type = FieldType.nullable(UuidType.INSTANCE); + try (final ListVector inVector = new ListVector("list", allocator, type, null)) { + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u2); + writer.endList(); + + writer.setValueCount(1); + + FieldReader reader = inVector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + Object result = inVector.getObject(0); + ArrayList resultSet = (ArrayList) result; + assertEquals(2, resultSet.size()); + assertEquals(u1, resultSet.get(0)); + assertEquals(u2, resultSet.get(1)); + } + } + + @Test + public void testListVectorReaderForExtensionType() throws Exception { + final FieldType type = FieldType.nullable(UuidType.INSTANCE); + try (final ListVector inVector = new ListVector("list", allocator, type, null)) { + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u2); + writer.endList(); + + writer.setValueCount(1); + + UnionListReader reader = inVector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(0); + reader.next(); + FieldReader uuidReader = reader.reader(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + } + } + + @Test + public void testCopyFromForExtensionType() throws Exception { + try (ListVector inVector = ListVector.empty("input", allocator); + ListVector outVector = ListVector.empty("output", allocator)) { + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startList(); + + writer.extension(UuidType.INSTANCE).writeExtension(u1); + writer.writeExtension(u2); + writer.writeNull(); + writer.endList(); + + writer.setValueCount(3); + + // copy values from input to output + outVector.allocateNew(); + outVector.copyFrom(0, 0, inVector); + outVector.setValueCount(3); + + UnionListReader reader = outVector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(0); + reader.next(); + FieldReader uuidReader = reader.reader(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + } + } + + @Test + public void testCopyValueSafeForExtensionType() throws Exception { + try (ListVector inVector = ListVector.empty("input", allocator); + ListVector outVector = ListVector.empty("output", allocator)) { + UnionListWriter writer = inVector.getWriter(); + writer.allocate(); + + // Create first list with UUIDs + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startList(); + ExtensionWriter extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1); + extensionWriter.writeExtension(u2); + writer.endList(); + + // Create second list with UUIDs + writer.setPosition(1); + UUID u3 = UUID.randomUUID(); + UUID u4 = UUID.randomUUID(); + writer.startList(); + extensionWriter = writer.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u3); + extensionWriter.writeExtension(u4); + extensionWriter.writeNull(); + + writer.endList(); + writer.setValueCount(2); + + // Use TransferPair with ExtensionTypeWriterFactory + // This tests the new makeTransferPair API with writerFactory parameter + outVector.allocateNew(); + TransferPair transferPair = inVector.makeTransferPair(outVector); + transferPair.copyValueSafe(0, 0); + transferPair.copyValueSafe(1, 1); + outVector.setValueCount(2); + + // Verify first list + UnionListReader reader = outVector.getReader(); + reader.setPosition(0); + assertTrue(reader.isSet(), "first list shouldn't be null"); + reader.next(); + FieldReader uuidReader = reader.reader(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + + // Verify second list + reader.setPosition(1); + assertTrue(reader.isSet(), "second list shouldn't be null"); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u3, actualUuid); + reader.next(); + uuidReader = reader.reader(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u4, actualUuid); + reader.next(); + uuidReader = reader.reader(); + assertFalse(uuidReader.isSet(), "third element should be null"); + } + } + + @Test + public void testEmptyListOffsetBuffer() { + // Test that ListVector has correct readableBytes after allocation. + // According to Arrow spec, offset buffer must have N+1 entries. + // Even when N=0, it should contain [0]. + try (ListVector list = ListVector.empty("list", allocator)) { + list.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); + list.allocateNew(); + list.setValueCount(0); + + List buffers = list.getFieldBuffers(); + assertTrue( + buffers.get(1).readableBytes() >= BaseRepeatedValueVector.OFFSET_WIDTH, + "Offset buffer should have at least " + + BaseRepeatedValueVector.OFFSET_WIDTH + + " bytes for offset[0]"); + assertEquals(0, list.getOffsetBuffer().getInt(0)); + } + } + private void writeIntValues(UnionListWriter writer, int[] values) { writer.startList(); for (int v : values) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java index 639585fc48..8ab0edb145 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestListViewVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -1075,7 +1076,7 @@ public void testGetBufferSizeFor() { int[] indices = new int[] {0, 2, 4, 6, 10, 14}; for (int valueCount = 1; valueCount <= 5; valueCount++) { - int validityBufferSize = BitVectorHelper.getValidityBufferSize(valueCount); + int validityBufferSize = getValidityBufferSizeFromCount(valueCount); int offsetBufferSize = valueCount * BaseRepeatedValueViewVector.OFFSET_WIDTH; int sizeBufferSize = valueCount * BaseRepeatedValueViewVector.SIZE_WIDTH; @@ -1549,55 +1550,7 @@ public void testOverwriteWithNull() { public void testOutOfOrderOffset1() { // [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] try (ListViewVector listViewVector = ListViewVector.empty("listview", allocator)) { - // Allocate buffers in listViewVector by calling `allocateNew` method. - listViewVector.allocateNew(); - - // Initialize the child vector using `initializeChildrenFromFields` method. - - FieldType fieldType = new FieldType(true, new ArrowType.Int(16, true), null, null); - Field field = new Field("child-vector", fieldType, null); - listViewVector.initializeChildrenFromFields(Collections.singletonList(field)); - - // Set values in the child vector. - FieldVector fieldVector = listViewVector.getDataVector(); - fieldVector.clear(); - - SmallIntVector childVector = (SmallIntVector) fieldVector; - - childVector.allocateNew(7); - - childVector.set(0, 0); - childVector.set(1, -127); - childVector.set(2, 127); - childVector.set(3, 50); - childVector.set(4, 12); - childVector.set(5, -7); - childVector.set(6, 25); - - childVector.setValueCount(7); - - // Set validity, offset and size buffers using `setValidity`, - // `setOffset` and `setSize` methods. - listViewVector.setValidity(0, 1); - listViewVector.setValidity(1, 0); - listViewVector.setValidity(2, 1); - listViewVector.setValidity(3, 1); - listViewVector.setValidity(4, 1); - - listViewVector.setOffset(0, 4); - listViewVector.setOffset(1, 7); - listViewVector.setOffset(2, 0); - listViewVector.setOffset(3, 0); - listViewVector.setOffset(4, 3); - - listViewVector.setSize(0, 3); - listViewVector.setSize(1, 0); - listViewVector.setSize(2, 4); - listViewVector.setSize(3, 0); - listViewVector.setSize(4, 2); - - // Set value count using `setValueCount` method. - listViewVector.setValueCount(5); + initializeListViewVectorAsInSpecification(listViewVector); final ArrowBuf offSetBuffer = listViewVector.getOffsetBuffer(); final ArrowBuf sizeBuffer = listViewVector.getSizeBuffer(); @@ -2216,6 +2169,105 @@ public void testRangeChildVector2() { } } + @Test + public void testGetElementStartIndexAndEndIndexOrderedOffsetsNoIntersection() { + /* + values = [10, 20, 30, 40, 50] + offsets = [0, 3] + sizes = [3, 2] + vector: [[10, 20, 30], [40, 50]] + */ + try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { + initializeListViewVector( + listViewVector, List.of(10, 20, 30, 40, 50), List.of(1, 1), List.of(0, 3), List.of(3, 2)); + + assertEquals(0, listViewVector.getElementStartIndex(0)); + assertEquals(3, listViewVector.getElementEndIndex(0)); + assertEquals(3, listViewVector.getElementStartIndex(1)); + assertEquals(5, listViewVector.getElementEndIndex(1)); + + final FieldVector dataVec = listViewVector.getDataVector(); + int elemIndex = 0; + int start = listViewVector.getElementStartIndex(elemIndex); + int end = listViewVector.getElementEndIndex(elemIndex); + List list = listViewVector.getObject(elemIndex); + assertEquals(end - start, list.size()); + for (int j = 0; j < list.size(); j++) { + assertEquals(((SmallIntVector) dataVec).get(start + j), list.get(j)); + } + } + } + + @Test + public void testGetElementStartIndexAndEndIndexNotOrderedOffsetsNoIntersection() { + /* + values = [1, 2, 3, 4, 5, 6] + validity = [1, 1, 1] + offsets = [4, 2, 0] + sizes = [2, 2, 2] + vector: [[5, 6], [3, 4], [1, 2]] + */ + try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { + initializeListViewVector( + listViewVector, + List.of(1, 2, 3, 4, 5, 6), + List.of(1, 1, 1), + List.of(4, 2, 0), + List.of(2, 2, 2)); + + assertEquals(4, listViewVector.getElementStartIndex(0)); + assertEquals(6, listViewVector.getElementEndIndex(0)); + assertEquals(2, listViewVector.getElementStartIndex(1)); + assertEquals(4, listViewVector.getElementEndIndex(1)); + assertEquals(0, listViewVector.getElementStartIndex(2)); + assertEquals(2, listViewVector.getElementEndIndex(2)); + } + } + + @Test + public void testGetElementStartIndexAndEndIndexOrderedOffsetsWithIntersection() { + /* + values = [1, 2, 3, 4, 5] + validity = [1, 1, 1] + offsets = [0, 1, 4] + sizes = [2, 3, 1] + vector: [[1, 2], [2, 3, 4], [5]] + */ + try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { + initializeListViewVector( + listViewVector, + List.of(1, 2, 3, 4, 5), + List.of(1, 1, 1), + List.of(0, 1, 4), + List.of(2, 3, 1)); + + assertEquals(0, listViewVector.getElementStartIndex(0)); + assertEquals(2, listViewVector.getElementEndIndex(0)); + assertEquals(1, listViewVector.getElementStartIndex(1)); + assertEquals(4, listViewVector.getElementEndIndex(1)); + assertEquals(4, listViewVector.getElementStartIndex(2)); + assertEquals(5, listViewVector.getElementEndIndex(2)); + } + } + + @Test + public void testGetElementStartIndexAndEndIndexOrderedOffsetsAsInSpecification() { + try (ListViewVector listViewVector = ListViewVector.empty("sourceVector", allocator)) { + initializeListViewVectorAsInSpecification(listViewVector); + + assertEquals(4, listViewVector.getElementStartIndex(0)); + assertEquals(7, listViewVector.getElementEndIndex(0)); + assertEquals(7, listViewVector.getElementStartIndex(1)); + assertEquals(7, listViewVector.getElementEndIndex(1)); + assertEquals(0, listViewVector.getElementStartIndex(2)); + assertEquals(4, listViewVector.getElementEndIndex(2)); + assertEquals(0, listViewVector.getElementStartIndex(3)); + assertEquals(0, listViewVector.getElementEndIndex(3)); + assertEquals(3, listViewVector.getElementStartIndex(4)); + assertEquals(5, listViewVector.getElementEndIndex(4)); + } + } + private void writeIntValues(UnionListViewWriter writer, int[] values) { writer.startListView(); for (int v : values) { @@ -2223,4 +2275,70 @@ private void writeIntValues(UnionListViewWriter writer, int[] values) { } writer.endListView(); } + + /** + * ListViewVector from the specification. + */ + private void initializeListViewVectorAsInSpecification(ListViewVector listViewVector) { + /* + values = [0, -127, 127, 50, 12, -7, 25] + validity = [1, 1, 1, 0, 1] (reversed) + offsets = [4, 7, 0, 0, 3] + sizes = [3, 0, 4, 0, 2] + vector: [[12, -7, 25], null, [0, -127, 127, 50], [], [50, 12]] + */ + initializeListViewVector( + listViewVector, + List.of(0, -127, 127, 50, 12, -7, 25), + List.of(1, 1, 1, 0, 1), + List.of(4, 7, 0, 0, 3), + List.of(3, 0, 4, 0, 2)); + } + + private void initializeListViewVector( + ListViewVector listViewVector, + List values, + List validity, + List offsets, + List sizes) { + // Allocate buffers in listViewVector by calling `allocateNew` method. + assert offsets.size() == sizes.size(); + listViewVector.allocateNew(); + + // Initialize the child vector using `initializeChildrenFromFields` method. + FieldType fieldType = new FieldType(true, new ArrowType.Int(16, true), null, null); + Field field = new Field("child-vector", fieldType, null); + listViewVector.initializeChildrenFromFields(Collections.singletonList(field)); + + // Set values in the child vector. + FieldVector fieldVector = listViewVector.getDataVector(); + fieldVector.clear(); + + SmallIntVector childVector = (SmallIntVector) fieldVector; + childVector.allocateNew(values.size()); + for (int i = 0; i < values.size(); i++) { + childVector.set(i, values.get(i)); + } + childVector.setValueCount(values.size()); + + // Set validity, offset and size buffers using `setValidity`, + // `setOffset` and `setSize` methods. + List reversedValidity = new ArrayList<>(validity); + Collections.reverse(reversedValidity); + for (int i = 0; i < reversedValidity.size(); i++) { + listViewVector.setValidity(i, reversedValidity.get(i)); + } + + for (int i = 0; i < offsets.size(); i++) { + listViewVector.setOffset(i, offsets.get(i)); + } + + for (int i = 0; i < sizes.size(); i++) { + listViewVector.setSize(i, sizes.get(i)); + } + + // Set value count using `setValueCount` method. + listViewVector.setValueCount(offsets.size()); + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java index 313d83ec91..2f520f3882 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestMapVector.java @@ -16,16 +16,19 @@ */ package org.apache.arrow.vector; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; 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 java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.MapVector; @@ -33,15 +36,20 @@ import org.apache.arrow.vector.complex.impl.UnionMapReader; import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.JsonStringArrayList; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -1263,4 +1271,395 @@ public void testMapTypeReturnsSupportedMapWriter() { assertEquals(11, getResultValue(resultStruct)); } } + + @Test + public void testMapVectorWithExtensionType() throws Exception { + try (final MapVector inVector = MapVector.empty("map", allocator, false)) { + inVector.allocateNew(); + UnionMapWriter writer = inVector.getWriter(); + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startMap(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(0); + ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1, UuidType.INSTANCE); + writer.endEntry(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(1); + extensionWriter = writer.value().extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u2, UuidType.INSTANCE); + writer.endEntry(); + writer.endMap(); + + writer.setValueCount(1); + + UnionMapReader mapReader = inVector.getReader(); + mapReader.setPosition(0); + mapReader.next(); + FieldReader uuidReader = mapReader.value(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + mapReader.next(); + uuidReader = mapReader.value(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + } + } + + @Test + public void testCopyFromForExtensionType() throws Exception { + try (final MapVector inVector = MapVector.empty("in", allocator, false); + final MapVector outVector = MapVector.empty("out", allocator, false)) { + inVector.allocateNew(); + UnionMapWriter writer = inVector.getWriter(); + writer.setPosition(0); + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + writer.startMap(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(0); + ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE); + extensionWriter.writeExtension(u1, UuidType.INSTANCE); + writer.endEntry(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(1); + extensionWriter.writeExtension(u2, UuidType.INSTANCE); + writer.endEntry(); + writer.endMap(); + + writer.setValueCount(1); + outVector.allocateNew(); + outVector.copyFrom(0, 0, inVector); + outVector.setValueCount(1); + + UnionMapReader mapReader = outVector.getReader(); + mapReader.setPosition(0); + mapReader.next(); + FieldReader uuidReader = mapReader.value(); + NullableUuidHolder holder = new NullableUuidHolder(); + uuidReader.read(holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u1, actualUuid); + mapReader.next(); + uuidReader = mapReader.value(); + uuidReader.read(holder); + actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(u2, actualUuid); + } + } + + /** + * Regression test for GH-586: UnionMapWriter.fixedSizeBinary() should properly delegate to the + * entry writer for both key and value paths. + */ + @Test + public void testFixedSizeBinaryWriter() { + try (MapVector mapVector = MapVector.empty("map_vector", allocator, false)) { + UnionMapWriter writer = mapVector.getWriter(); + writer.allocate(); + + // populate input vector with the following records + // {[11, 22] -> [32, 21]} + // {1 -> [11, 22], 2 -> [32, 21]} + // null + // {[11, 22] -> 1, [32, 21] -> 2} + // {[11, 22] -> null} + // {null -> [32, 21]} - wrong "for a given entry, the "key" is non-nullable" - todo: it + // shouldn't work. Should it? + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + + writer.setPosition(0); // optional + writer.startMap(); + writer.startEntry(); + writer + .key() + .fixedSizeBinary(holder1.byteWidth) + .write(holder1); // need to initialize with byteWidth - NPE otherwise + writer.value().fixedSizeBinary(holder2.byteWidth).write(holder2); + writer.endEntry(); + holder1.buffer.close(); + holder2.buffer.close(); + writer.endMap(); + + // {1 -> [11, 22], 2 -> [32, 21]} + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + writer.setPosition(1); + writer.startMap(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(1); + writer.value().fixedSizeBinary().write(holder1); + writer.endEntry(); + holder1.buffer.close(); + writer.startEntry(); + writer.key().bigInt().writeBigInt(2); + writer.value().fixedSizeBinary().write(holder2); + writer.endEntry(); + writer.endMap(); + holder2.buffer.close(); + + // {[11, 22] -> 1, [32, 21] -> 2} + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + writer.setPosition(3); + writer.startMap(); + writer.startEntry(); + writer.key().fixedSizeBinary().write(holder1); + writer.value().bigInt().writeBigInt(1); + writer.endEntry(); + holder1.buffer.close(); + writer.startEntry(); + writer.key().fixedSizeBinary().write(holder2); + writer.value().bigInt().writeBigInt(2); + writer.endEntry(); + writer.endMap(); + holder2.buffer.close(); + + // {[11, 22] -> null} + holder1 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + writer.setPosition(4); + writer.startMap(); + writer.startEntry(); + writer.key().fixedSizeBinary().write(holder1); + writer.endEntry(); + writer.endMap(); + holder1.buffer.close(); + + // {null -> [32, 21]} + holder2 = TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + writer.setPosition(5); + writer.startMap(); + writer.startEntry(); + writer.value().fixedSizeBinary().write(holder2); + writer.endEntry(); + writer.endMap(); + holder2.buffer.close(); + + writer.setValueCount(6); + + // assert the output vector is correct + FieldReader reader = mapVector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(1); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(2); + assertFalse(reader.isSet(), "should be null"); + reader.setPosition(3); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(4); + assertTrue(reader.isSet(), "shouldn't be null"); + reader.setPosition(5); + assertTrue(reader.isSet(), "shouldn't be null"); + + /* index 0 */ + Object result = mapVector.getObject(0); + ArrayList resultSet = (ArrayList) result; + assertEquals(1, resultSet.size()); + Map resultStruct = (Map) resultSet.get(0); + assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); + assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); + assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME)); + assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME)); + + /* index 1 */ + result = mapVector.getObject(1); + resultSet = (ArrayList) result; + assertEquals(2, resultSet.size()); + resultStruct = (Map) resultSet.get(0); + assertEquals(1L, getResultKey(resultStruct)); + assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); + assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.VALUE_NAME)); + resultStruct = (Map) resultSet.get(1); + assertEquals(2L, getResultKey(resultStruct)); + assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); + assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME)); + + /* index 2 */ + result = mapVector.getObject(2); + assertNull(result); + + /* index 3 */ + result = mapVector.getObject(3); + resultSet = (ArrayList) result; + assertEquals(2, resultSet.size()); + resultStruct = (Map) resultSet.get(0); + assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); + assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME)); + assertEquals(1L, getResultValue(resultStruct)); + resultStruct = (Map) resultSet.get(1); + assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); + assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.KEY_NAME)); + assertEquals(2L, getResultValue(resultStruct)); + + /* index 4 */ + result = mapVector.getObject(4); + resultSet = (ArrayList) result; + assertEquals(1, resultSet.size()); + resultStruct = (Map) resultSet.get(0); + assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); + assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME)); + assertFalse(resultStruct.containsKey(MapVector.VALUE_NAME)); + + /* index 5 */ + result = mapVector.getObject(5); + resultSet = (ArrayList) result; + assertEquals(1, resultSet.size()); + resultStruct = (Map) resultSet.get(0); + assertFalse(resultStruct.containsKey(MapVector.KEY_NAME)); + assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); + assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME)); + } + } + + @Test + public void testFixedSizeBinaryFirstInitialization() { + try (MapVector mapVector = MapVector.empty("map_vector", allocator, false)) { + UnionMapWriter writer = mapVector.getWriter(); + writer.allocate(); + + // populate input vector with the following records + // {[11, 22] -> [32, 21]} + FixedSizeBinaryHolder holder1 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {11, 22}); + FixedSizeBinaryHolder holder2 = + TestUtils.fixedSizeBinaryHolder(allocator, new byte[] {32, 21}); + + writer.setPosition(0); // optional + writer.startMap(); + writer.startEntry(); + // require byteWidth parameter for first-time initialization of `key` or `value` writers + assertThrows(NullPointerException.class, () -> writer.key().fixedSizeBinary().write(holder1)); + assertThrows( + NullPointerException.class, () -> writer.value().fixedSizeBinary().write(holder2)); + writer.key().fixedSizeBinary(holder1.byteWidth).write(holder1); + writer.value().fixedSizeBinary(holder2.byteWidth).write(holder2); + writer.endEntry(); + holder1.buffer.close(); + holder2.buffer.close(); + writer.endMap(); + + writer.setValueCount(1); + + // assert the output vector is correct + FieldReader reader = mapVector.getReader(); + assertTrue(reader.isSet(), "shouldn't be null"); + + /* index 0 */ + Object result = mapVector.getObject(0); + ArrayList resultSet = (ArrayList) result; + assertEquals(1, resultSet.size()); + Map resultStruct = (Map) resultSet.get(0); + assertTrue(resultStruct.containsKey(MapVector.KEY_NAME)); + assertTrue(resultStruct.containsKey(MapVector.VALUE_NAME)); + assertArrayEquals(new byte[] {11, 22}, (byte[]) resultStruct.get(MapVector.KEY_NAME)); + assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME)); + } + } + + @Test + public void testMapWithUuidKeyAndListUuidValue() throws Exception { + try (final MapVector mapVector = MapVector.empty("map", allocator, false)) { + mapVector.allocateNew(); + UnionMapWriter writer = mapVector.getWriter(); + + // Create test UUIDs + UUID key1 = UUID.randomUUID(); + UUID key2 = UUID.randomUUID(); + UUID value1a = UUID.randomUUID(); + UUID value1b = UUID.randomUUID(); + UUID value2a = UUID.randomUUID(); + UUID value2b = UUID.randomUUID(); + UUID value2c = UUID.randomUUID(); + + // Write first map entry: {key1 -> [value1a, value1b]} + writer.setPosition(0); + writer.startMap(); + + writer.startEntry(); + ExtensionWriter keyWriter = writer.key().extension(UuidType.INSTANCE); + keyWriter.writeExtension(key1, UuidType.INSTANCE); + ListWriter valueWriter = writer.value().list(); + valueWriter.startList(); + ExtensionWriter listItemWriter = valueWriter.extension(UuidType.INSTANCE); + listItemWriter.writeExtension(value1a, UuidType.INSTANCE); + listItemWriter = valueWriter.extension(UuidType.INSTANCE); + listItemWriter.writeExtension(value1b, UuidType.INSTANCE); + valueWriter.endList(); + writer.endEntry(); + + writer.startEntry(); + keyWriter = writer.key().extension(UuidType.INSTANCE); + keyWriter.writeExtension(key2, UuidType.INSTANCE); + valueWriter = writer.value().list(); + valueWriter.startList(); + listItemWriter = valueWriter.extension(UuidType.INSTANCE); + listItemWriter.writeExtension(value2a, UuidType.INSTANCE); + listItemWriter = valueWriter.extension(UuidType.INSTANCE); + listItemWriter.writeExtension(value2b, UuidType.INSTANCE); + listItemWriter = valueWriter.extension(UuidType.INSTANCE); + listItemWriter.writeExtension(value2c, UuidType.INSTANCE); + valueWriter.endList(); + writer.endEntry(); + + writer.endMap(); + writer.setValueCount(1); + + // Read and verify the data + UnionMapReader mapReader = mapVector.getReader(); + mapReader.setPosition(0); + + // Read first entry + mapReader.next(); + FieldReader keyReader = mapReader.key(); + NullableUuidHolder keyHolder = new NullableUuidHolder(); + keyReader.read(keyHolder); + UUID actualKey = UuidUtility.uuidFromArrowBuf(keyHolder.buffer, keyHolder.start); + assertEquals(key1, actualKey); + + FieldReader valueReader = mapReader.value(); + assertTrue(valueReader.isSet()); + List listValue = (List) valueReader.readObject(); + assertEquals(2, listValue.size()); + + // Verify first list item - readObject() returns UUID objects for extension types + UUID actualValue1a = (UUID) listValue.get(0); + assertEquals(value1a, actualValue1a); + + // Verify second list item + UUID actualValue1b = (UUID) listValue.get(1); + assertEquals(value1b, actualValue1b); + + // Read second entry + mapReader.next(); + keyReader = mapReader.key(); + keyReader.read(keyHolder); + actualKey = UuidUtility.uuidFromArrowBuf(keyHolder.buffer, keyHolder.start); + assertEquals(key2, actualKey); + + valueReader = mapReader.value(); + assertTrue(valueReader.isSet()); + listValue = (List) valueReader.readObject(); + assertEquals(3, listValue.size()); + + // Verify first list item - readObject() returns UUID objects for extension types + UUID actualValue2a = (UUID) listValue.get(0); + assertEquals(value2a, actualValue2a); + + // Verify second list item + UUID actualValue2b = (UUID) listValue.get(1); + assertEquals(value2b, actualValue2b); + + // Verify third list item + UUID actualValue2c = (UUID) listValue.get(2); + assertEquals(value2c, actualValue2c); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java b/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java index adf51c0730..9fa153e928 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestRunEndEncodedVector.java @@ -148,12 +148,18 @@ public void testRangeCompare() { assertTrue( constantVector.accept( new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 2, 13))); - assertFalse( - constantVector.accept( - new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 10, 10))); - assertFalse( - constantVector.accept( - new RangeEqualsVisitor(constantVector, constantVector), new Range(10, 1, 10))); + + // throws exception if the range end is out the bound of the vector + assertThrows( + IllegalArgumentException.class, + () -> + constantVector.accept( + new RangeEqualsVisitor(constantVector, constantVector), new Range(1, 10, 10))); + assertThrows( + IllegalArgumentException.class, + () -> + constantVector.accept( + new RangeEqualsVisitor(constantVector, constantVector), new Range(10, 1, 10))); // Create REE vector representing: [1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5]. RunEndEncodedVector reeVector = diff --git a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java index d40af9ae89..8c8a45f588 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java @@ -35,6 +35,7 @@ import org.apache.arrow.vector.complex.impl.NullableStructWriter; import org.apache.arrow.vector.complex.writer.Float8Writer; import org.apache.arrow.vector.complex.writer.IntWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.ComplexHolder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.Types.MinorType; @@ -42,7 +43,6 @@ import org.apache.arrow.vector.types.pojo.ArrowType.Struct; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.UuidType; import org.apache.arrow.vector.util.TransferPair; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -160,17 +160,23 @@ public void testGetPrimitiveVectors() { UnionVector unionVector = vector.addOrGetUnion("union"); unionVector.addVector(new BigIntVector("bigInt", allocator)); unionVector.addVector(new SmallIntVector("smallInt", allocator)); + unionVector.addVector(new UuidVector("uuid", allocator)); // add varchar vector vector.addOrGet( "varchar", FieldType.nullable(MinorType.VARCHAR.getType()), VarCharVector.class); + // add extension vector + vector.addOrGet("extension", FieldType.nullable(UuidType.INSTANCE), UuidVector.class); + List primitiveVectors = vector.getPrimitiveVectors(); - assertEquals(4, primitiveVectors.size()); + assertEquals(6, primitiveVectors.size()); assertEquals(MinorType.INT, primitiveVectors.get(0).getMinorType()); assertEquals(MinorType.BIGINT, primitiveVectors.get(1).getMinorType()); assertEquals(MinorType.SMALLINT, primitiveVectors.get(2).getMinorType()); - assertEquals(MinorType.VARCHAR, primitiveVectors.get(3).getMinorType()); + assertEquals(MinorType.EXTENSIONTYPE, primitiveVectors.get(3).getMinorType()); + assertEquals(MinorType.VARCHAR, primitiveVectors.get(4).getMinorType()); + assertEquals(MinorType.EXTENSIONTYPE, primitiveVectors.get(5).getMinorType()); } } @@ -341,7 +347,7 @@ public void testGetTransferPairWithFieldAndCallBack() { @Test public void testStructVectorWithExtensionTypes() { - UuidType uuidType = new UuidType(); + UuidType uuidType = UuidType.INSTANCE; Field uuidField = new Field("struct_child", FieldType.nullable(uuidType), null); Field structField = new Field("struct", FieldType.nullable(new ArrowType.Struct()), List.of(uuidField)); @@ -353,7 +359,7 @@ public void testStructVectorWithExtensionTypes() { @Test public void testStructVectorTransferPairWithExtensionType() { - UuidType uuidType = new UuidType(); + UuidType uuidType = UuidType.INSTANCE; Field uuidField = new Field("uuid_child", FieldType.nullable(uuidType), null); Field structField = new Field("struct", FieldType.nullable(new ArrowType.Struct()), List.of(uuidField)); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java index 82295f8037..d91b2004c0 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java @@ -18,8 +18,10 @@ import java.util.Random; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry; import org.apache.arrow.vector.types.pojo.FieldType; public class TestUtils { @@ -62,4 +64,26 @@ public static String generateRandomString(int length) { } return sb.toString(); } + + /* + * Ensure the extension type is registered, as there might other tests trying to unregister the + * type. ex.: TestExtensionType#readUnderlyingType + */ + public static void ensureRegistered(ArrowType.ExtensionType type) { + if (ExtensionTypeRegistry.lookup(type.extensionName()) == null) { + ExtensionTypeRegistry.register(type); + } + } + + public static FixedSizeBinaryHolder fixedSizeBinaryHolder( + BufferAllocator allocator, byte[] array) { + FixedSizeBinaryHolder holder = new FixedSizeBinaryHolder(); + holder.byteWidth = array.length; + holder.buffer = allocator.buffer(array.length); + for (int i = 0; i < array.length; i++) { + holder.buffer.setByte(i, array[i]); + } + + return holder; + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java b/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java new file mode 100644 index 0000000000..99045d1cba --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestUuidType.java @@ -0,0 +1,276 @@ +/* + * 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.arrow.vector; + +import static org.apache.arrow.vector.TestUtils.ensureRegistered; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +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 java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.UUID; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.dictionary.DictionaryProvider; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.UuidUtility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestUuidType { + BufferAllocator allocator; + + @BeforeEach + void beforeEach() { + allocator = new RootAllocator(); + } + + @AfterEach + void afterEach() { + allocator.close(); + } + + @Test + void testConstants() { + assertEquals("arrow.uuid", UuidType.EXTENSION_NAME); + assertNotNull(UuidType.INSTANCE); + assertNotNull(UuidType.STORAGE_TYPE); + assertInstanceOf(ArrowType.FixedSizeBinary.class, UuidType.STORAGE_TYPE); + assertEquals( + UuidType.UUID_BYTE_WIDTH, + ((ArrowType.FixedSizeBinary) UuidType.STORAGE_TYPE).getByteWidth()); + } + + @Test + void testStorageType() { + UuidType type = UuidType.INSTANCE; + assertEquals(UuidType.STORAGE_TYPE, type.storageType()); + assertInstanceOf(ArrowType.FixedSizeBinary.class, type.storageType()); + } + + @Test + void testExtensionName() { + UuidType type = UuidType.INSTANCE; + assertEquals("arrow.uuid", type.extensionName()); + } + + @Test + void testExtensionEquals() { + UuidType type1 = UuidType.INSTANCE; + UuidType type2 = UuidType.INSTANCE; + UuidType type3 = UuidType.INSTANCE; + + assertTrue(type1.extensionEquals(type2)); + assertTrue(type1.extensionEquals(type3)); + assertTrue(type2.extensionEquals(type3)); + } + + @Test + void testIsComplex() { + UuidType type = UuidType.INSTANCE; + assertFalse(type.isComplex()); + } + + @Test + void testSerialize() { + UuidType type = UuidType.INSTANCE; + String serialized = type.serialize(); + assertEquals("", serialized); + } + + @Test + void testDeserializeValid() { + UuidType type = UuidType.INSTANCE; + ArrowType storageType = new ArrowType.FixedSizeBinary(UuidType.UUID_BYTE_WIDTH); + + ArrowType deserialized = assertDoesNotThrow(() -> type.deserialize(storageType, "")); + assertInstanceOf(UuidType.class, deserialized); + assertEquals(UuidType.INSTANCE, deserialized); + } + + @Test + void testDeserializeInvalidStorageType() { + UuidType type = UuidType.INSTANCE; + ArrowType wrongStorageType = new ArrowType.FixedSizeBinary(32); + + assertThrows(UnsupportedOperationException.class, () -> type.deserialize(wrongStorageType, "")); + } + + @Test + void testGetNewVector() { + UuidType type = UuidType.INSTANCE; + try (FieldVector vector = + type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) { + assertInstanceOf(UuidVector.class, vector); + assertEquals("uuid_field", vector.getField().getName()); + assertEquals(type, vector.getField().getType()); + } + } + + @Test + void testVectorOperations() { + UuidType type = UuidType.INSTANCE; + try (FieldVector vector = + type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) { + UuidVector uuidVector = (UuidVector) vector; + + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + uuidVector.setSafe(0, uuid1); + uuidVector.setSafe(1, uuid2); + uuidVector.setNull(2); + uuidVector.setValueCount(3); + + assertEquals(uuid1, uuidVector.getObject(0)); + assertEquals(uuid2, uuidVector.getObject(1)); + assertNull(uuidVector.getObject(2)); + assertFalse(uuidVector.isNull(0)); + assertFalse(uuidVector.isNull(1)); + assertTrue(uuidVector.isNull(2)); + } + } + + @Test + void testIpcRoundTrip() { + UuidType type = UuidType.INSTANCE; + ensureRegistered(type); + + Schema schema = new Schema(Collections.singletonList(Field.nullable("uuid", type))); + byte[] serialized = schema.serializeAsMessage(); + Schema deserialized = Schema.deserializeMessage(ByteBuffer.wrap(serialized)); + assertEquals(schema, deserialized); + } + + @Test + void testVectorIpcRoundTrip() throws IOException { + UuidType type = UuidType.INSTANCE; + ensureRegistered(type); + + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + try (FieldVector vector = type.getNewVector("field", FieldType.nullable(type), allocator)) { + UuidVector uuidVector = (UuidVector) vector; + uuidVector.setSafe(0, uuid1); + uuidVector.setNull(1); + uuidVector.setSafe(2, uuid2); + uuidVector.setValueCount(3); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (VectorSchemaRoot root = new VectorSchemaRoot(Collections.singletonList(uuidVector)); + ArrowStreamWriter writer = + new ArrowStreamWriter(root, new DictionaryProvider.MapDictionaryProvider(), baos)) { + writer.start(); + writer.writeBatch(); + } + + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(baos.toByteArray()), allocator)) { + assertTrue(reader.loadNextBatch()); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertEquals(3, root.getRowCount()); + assertEquals( + new Schema(Collections.singletonList(uuidVector.getField())), root.getSchema()); + + UuidVector actual = assertInstanceOf(UuidVector.class, root.getVector("field")); + assertFalse(actual.isNull(0)); + assertTrue(actual.isNull(1)); + assertFalse(actual.isNull(2)); + assertEquals(uuid1, actual.getObject(0)); + assertNull(actual.getObject(1)); + assertEquals(uuid2, actual.getObject(2)); + } + } + } + + @Test + void testVectorByteArrayOperations() { + UuidType type = UuidType.INSTANCE; + try (FieldVector vector = + type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) { + UuidVector uuidVector = (UuidVector) vector; + + UUID uuid = UUID.randomUUID(); + byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid); + + uuidVector.setSafe(0, uuidBytes); + uuidVector.setValueCount(1); + + assertEquals(uuid, uuidVector.getObject(0)); + + // Verify the bytes match + byte[] actualBytes = new byte[UuidType.UUID_BYTE_WIDTH]; + int offset = uuidVector.getStartOffset(0); + uuidVector.getDataBuffer().getBytes(offset, actualBytes); + assertArrayEquals(uuidBytes, actualBytes); + } + } + + @Test + void testGetNewVectorWithCustomFieldType() { + UuidType type = UuidType.INSTANCE; + FieldType fieldType = new FieldType(false, type, null); + + try (FieldVector vector = type.getNewVector("non_nullable_uuid", fieldType, allocator)) { + assertInstanceOf(UuidVector.class, vector); + assertEquals("non_nullable_uuid", vector.getField().getName()); + assertFalse(vector.getField().isNullable()); + } + } + + @Test + void testSingleton() { + UuidType type1 = UuidType.INSTANCE; + UuidType type2 = UuidType.INSTANCE; + + // Same instance + assertSame(type1, type2); + assertTrue(type1.extensionEquals(type2)); + } + + @Test + void testUnderlyingVector() { + UuidType type = UuidType.INSTANCE; + try (FieldVector vector = + type.getNewVector("uuid_field", FieldType.nullable(type), allocator)) { + UuidVector uuidVector = (UuidVector) vector; + FixedSizeBinaryVector underlying = uuidVector.getUnderlyingVector(); + + assertInstanceOf(FixedSizeBinaryVector.class, underlying); + assertEquals(UuidType.UUID_BYTE_WIDTH, underlying.getByteWidth()); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java new file mode 100644 index 0000000000..b5dd12d89c --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestUuidVector.java @@ -0,0 +1,726 @@ +/* + * 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.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +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 java.nio.ByteBuffer; +import java.util.UUID; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.complex.impl.NullableUuidHolderReaderImpl; +import org.apache.arrow.vector.complex.impl.UuidReaderImpl; +import org.apache.arrow.vector.complex.impl.UuidWriterImpl; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; +import org.apache.arrow.vector.holders.UuidHolder; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.util.UuidUtility; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for UuidVector, UuidWriterImpl, and UuidReaderImpl. */ +class TestUuidVector { + + private BufferAllocator allocator; + + @BeforeEach + void beforeEach() { + allocator = new RootAllocator(); + } + + @AfterEach + void afterEach() { + allocator.close(); + } + + // ========== Writer Tests ========== + + @Test + void testWriteToExtensionVector() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + UUID uuid = UUID.randomUUID(); + ByteBuffer bb = ByteBuffer.allocate(UuidType.UUID_BYTE_WIDTH); + bb.putLong(uuid.getMostSignificantBits()); + bb.putLong(uuid.getLeastSignificantBits()); + + // Allocate ArrowBuf for the holder + try (ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) { + buf.setBytes(0, bb.array()); + + UuidHolder holder = new UuidHolder(); + holder.buffer = buf; + + writer.write(holder); + UUID result = vector.getObject(0); + assertEquals(uuid, result); + } + } + } + + @Test + void testWriteExtensionWithUUID() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + UUID uuid = UUID.randomUUID(); + writer.setPosition(0); + writer.writeExtension(uuid); + + UUID result = vector.getObject(0); + assertEquals(uuid, result); + assertEquals(1, vector.getValueCount()); + } + } + + @Test + void testWriteExtensionWithByteArray() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + UUID uuid = UUID.randomUUID(); + byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid); + + writer.setPosition(0); + writer.writeExtension(uuidBytes); + + UUID result = vector.getObject(0); + assertEquals(uuid, result); + assertEquals(1, vector.getValueCount()); + } + } + + @Test + void testWriteExtensionWithArrowBuf() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector); + ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) { + UUID uuid = UUID.randomUUID(); + byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid); + buf.setBytes(0, uuidBytes); + + writer.setPosition(0); + writer.writeExtension(buf); + + UUID result = vector.getObject(0); + assertEquals(uuid, result); + assertEquals(1, vector.getValueCount()); + } + } + + @Test + void testWriteExtensionWithUnsupportedType() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + writer.setPosition(0); + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> writer.writeExtension("invalid-type")); + + assertTrue( + exception.getMessage().contains("Unsupported value type for UUID: java.lang.String")); + } + } + + @Test + void testWriteExtensionMultipleValues() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + UUID uuid3 = UUID.randomUUID(); + + writer.setPosition(0); + writer.writeExtension(uuid1); + writer.setPosition(1); + writer.writeExtension(uuid2); + writer.setPosition(2); + writer.writeExtension(uuid3); + + assertEquals(uuid1, vector.getObject(0)); + assertEquals(uuid2, vector.getObject(1)); + assertEquals(uuid3, vector.getObject(2)); + assertEquals(3, vector.getValueCount()); + } + } + + @Test + void testWriteWithUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector); + ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) { + UUID uuid = UUID.randomUUID(); + byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid); + buf.setBytes(0, uuidBytes); + + UuidHolder holder = new UuidHolder(); + holder.buffer = buf; + holder.isSet = 1; + + writer.setPosition(0); + writer.write(holder); + + UUID result = vector.getObject(0); + assertEquals(uuid, result); + assertEquals(1, vector.getValueCount()); + } + } + + @Test + void testWriteWithNullableUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector); + ArrowBuf buf = allocator.buffer(UuidType.UUID_BYTE_WIDTH)) { + UUID uuid = UUID.randomUUID(); + byte[] uuidBytes = UuidUtility.getBytesFromUUID(uuid); + buf.setBytes(0, uuidBytes); + + NullableUuidHolder holder = new NullableUuidHolder(); + holder.buffer = buf; + holder.isSet = 1; + + writer.setPosition(0); + writer.write(holder); + + UUID result = vector.getObject(0); + assertEquals(uuid, result); + assertEquals(1, vector.getValueCount()); + } + } + + @Test + void testWriteWithNullableUuidHolderNull() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + NullableUuidHolder holder = new NullableUuidHolder(); + holder.isSet = 0; + + writer.setPosition(0); + writer.write(holder); + + assertTrue(vector.isNull(0)); + assertEquals(1, vector.getValueCount()); + } + } + + // ========== Reader Tests ========== + + @Test + void testReaderCopyAsValueExtensionVector() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator); + UuidVector vectorForRead = new UuidVector("test2", allocator); + UuidWriterImpl writer = new UuidWriterImpl(vector)) { + UUID uuid = UUID.randomUUID(); + vectorForRead.setValueCount(1); + vectorForRead.set(0, uuid); + UuidReaderImpl reader = (UuidReaderImpl) vectorForRead.getReader(); + reader.copyAsValue(writer); + UuidReaderImpl reader2 = (UuidReaderImpl) vector.getReader(); + NullableUuidHolder holder = new NullableUuidHolder(); + reader2.read(0, holder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(uuid, actualUuid); + } + } + + @Test + void testReaderReadWithUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setValueCount(1); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + reader.setPosition(0); + + NullableUuidHolder holder = new NullableUuidHolder(); + reader.read(holder); + + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(uuid, actualUuid); + assertEquals(1, holder.isSet); + } + } + + @Test + void testReaderReadWithNullableUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setValueCount(1); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + reader.setPosition(0); + + NullableUuidHolder holder = new NullableUuidHolder(); + reader.read(holder); + + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(uuid, actualUuid); + assertEquals(1, holder.isSet); + } + } + + @Test + void testReaderReadWithNullableUuidHolderNull() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + vector.setNull(0); + vector.setValueCount(1); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + reader.setPosition(0); + + NullableUuidHolder holder = new NullableUuidHolder(); + reader.read(holder); + + assertEquals(0, holder.isSet); + } + } + + @Test + void testReaderReadWithArrayIndexUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + UUID uuid3 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setSafe(1, uuid2); + vector.setSafe(2, uuid3); + vector.setValueCount(3); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + + NullableUuidHolder holder = new NullableUuidHolder(); + reader.read(1, holder); + + UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start); + assertEquals(uuid2, actualUuid); + assertEquals(1, holder.isSet); + } + } + + @Test + void testReaderReadWithArrayIndexNullableUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setNull(1); + vector.setSafe(2, uuid2); + vector.setValueCount(3); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + + NullableUuidHolder holder1 = new NullableUuidHolder(); + reader.read(0, holder1); + assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder1.buffer, holder1.start)); + assertEquals(1, holder1.isSet); + + NullableUuidHolder holder2 = new NullableUuidHolder(); + reader.read(1, holder2); + assertEquals(0, holder2.isSet); + + NullableUuidHolder holder3 = new NullableUuidHolder(); + reader.read(2, holder3); + assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder3.buffer, holder3.start)); + assertEquals(1, holder3.isSet); + } + } + + @Test + void testReaderReadWithUnsupportedHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setValueCount(1); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + reader.setPosition(0); + + // Create a mock unsupported holder + ExtensionHolder unsupportedHolder = + new ExtensionHolder() { + @Override + public ArrowType type() { + return null; + } + }; + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> reader.read(unsupportedHolder)); + + assertTrue(exception.getMessage().contains("Unsupported holder type for UuidReader")); + } + } + + @Test + void testReaderIsSet() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setNull(1); + vector.setSafe(2, uuid); + vector.setValueCount(3); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + + reader.setPosition(0); + assertTrue(reader.isSet()); + + reader.setPosition(1); + assertFalse(reader.isSet()); + + reader.setPosition(2); + assertTrue(reader.isSet()); + } + } + + @Test + void testReaderReadObject() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setNull(1); + vector.setSafe(2, uuid2); + vector.setValueCount(3); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + + reader.setPosition(0); + assertEquals(uuid1, reader.readObject()); + + reader.setPosition(1); + assertNull(reader.readObject()); + + reader.setPosition(2); + assertEquals(uuid2, reader.readObject()); + } + } + + @Test + void testReaderGetMinorType() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + assertEquals(vector.getMinorType(), reader.getMinorType()); + } + } + + @Test + void testReaderGetField() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + assertEquals(vector.getField(), reader.getField()); + assertEquals("test", reader.getField().getName()); + } + } + + @Test + void testHolderStartOffsetWithMultipleValues() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + UUID uuid3 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setSafe(1, uuid2); + vector.setSafe(2, uuid3); + vector.setValueCount(3); + + // Test UuidHolder with different indices + NullableUuidHolder holder = new NullableUuidHolder(); + vector.get(0, holder); + assertEquals(0, holder.start); + assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + + vector.get(1, holder); + assertEquals(16, holder.start); // UUID_BYTE_WIDTH = 16 + assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + + vector.get(2, holder); + assertEquals(32, holder.start); // 2 * UUID_BYTE_WIDTH = 32 + assertEquals(uuid3, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + } + } + + @Test + void testNullableHolderStartOffsetWithMultipleValues() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setNull(1); + vector.setSafe(2, uuid2); + vector.setValueCount(3); + + // Test NullableUuidHolder with different indices + NullableUuidHolder holder1 = new NullableUuidHolder(); + vector.get(0, holder1); + assertEquals(0, holder1.start); + assertEquals(1, holder1.isSet); + assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder1.buffer, holder1.start)); + + NullableUuidHolder holder2 = new NullableUuidHolder(); + vector.get(1, holder2); + assertEquals(0, holder2.isSet); + + NullableUuidHolder holder3 = new NullableUuidHolder(); + vector.get(2, holder3); + assertEquals(32, holder3.start); // 2 * UUID_BYTE_WIDTH = 32 + assertEquals(1, holder3.isSet); + assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder3.buffer, holder3.start)); + + // Verify all holders share the same buffer + assertEquals(holder1.buffer, holder3.buffer); + } + } + + @Test + void testSetFromHolderWithStartOffset() throws Exception { + try (UuidVector sourceVector = new UuidVector("source", allocator); + UuidVector targetVector = new UuidVector("target", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + sourceVector.setSafe(0, uuid1); + sourceVector.setSafe(1, uuid2); + sourceVector.setValueCount(3); + + // Get holder from index 1 (should have start = 16) + NullableUuidHolder holder = new NullableUuidHolder(); + sourceVector.get(1, holder); + assertEquals(16, holder.start); + + // Set target vector using holder with non-zero start offset + targetVector.setSafe(0, holder); + targetVector.setValueCount(1); + + // Verify the value was copied correctly + assertEquals(uuid2, targetVector.getObject(0)); + } + } + + @Test + void testSetFromNullableHolderWithStartOffset() throws Exception { + try (UuidVector sourceVector = new UuidVector("source", allocator); + UuidVector targetVector = new UuidVector("target", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + + sourceVector.setSafe(0, uuid1); + sourceVector.setNull(1); + sourceVector.setSafe(2, uuid2); + sourceVector.setValueCount(3); + + // Get holder from index 2 (should have start = 32) + NullableUuidHolder holder = new NullableUuidHolder(); + sourceVector.get(2, holder); + assertEquals(32, holder.start); + assertEquals(1, holder.isSet); + + // Set target vector using holder with non-zero start offset + targetVector.setSafe(0, holder); + targetVector.setValueCount(1); + + // Verify the value was copied correctly + assertEquals(uuid2, targetVector.getObject(0)); + + // Test with null holder + NullableUuidHolder nullHolder = new NullableUuidHolder(); + sourceVector.get(1, nullHolder); + assertEquals(0, nullHolder.isSet); + + targetVector.setSafe(1, nullHolder); + targetVector.setValueCount(2); + assertTrue(targetVector.isNull(1)); + } + } + + @Test + void testGetStartOffset() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + vector.allocateNew(10); + + // Test getStartOffset for various indices + assertEquals(0, vector.getStartOffset(0)); + assertEquals(16, vector.getStartOffset(1)); + assertEquals(32, vector.getStartOffset(2)); + assertEquals(48, vector.getStartOffset(3)); + assertEquals(160, vector.getStartOffset(10)); + } + } + + @Test + void testReaderWithStartOffsetMultipleReads() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + UUID uuid3 = UUID.randomUUID(); + + vector.setSafe(0, uuid1); + vector.setSafe(1, uuid2); + vector.setSafe(2, uuid3); + vector.setValueCount(3); + + UuidReaderImpl reader = (UuidReaderImpl) vector.getReader(); + NullableUuidHolder holder = new NullableUuidHolder(); + + // Read from different positions and verify start offset + reader.read(0, holder); + assertEquals(0, holder.start); + assertEquals(uuid1, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + + reader.read(1, holder); + assertEquals(16, holder.start); + assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + + reader.read(2, holder); + assertEquals(32, holder.start); + assertEquals(uuid3, UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start)); + } + } + + @Test + void testWriterWithExtensionHolder() throws Exception { + try (UuidVector sourceVector = new UuidVector("source", allocator); + UuidVector targetVector = new UuidVector("target", allocator)) { + UUID uuid = UUID.randomUUID(); + sourceVector.setSafe(0, uuid); + sourceVector.setValueCount(1); + + // Get holder from source + NullableUuidHolder holder = new NullableUuidHolder(); + sourceVector.get(0, holder); + + // Write using UuidWriterImpl with ExtensionHolder + UuidWriterImpl writer = new UuidWriterImpl(targetVector); + writer.setPosition(0); + writer.writeExtension(holder); + + assertEquals(uuid, targetVector.getObject(0)); + } + } + + @Test + void testNullableUuidHolderReaderImpl() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setValueCount(1); + + // Get holder from vector + NullableUuidHolder sourceHolder = new NullableUuidHolder(); + vector.get(0, sourceHolder); + assertEquals(1, sourceHolder.isSet); + assertEquals(0, sourceHolder.start); + + // Create reader from holder + NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder); + assertTrue(reader.isSet()); + assertEquals(uuid, reader.readObject()); + + // Read into another holder + NullableUuidHolder targetHolder = new NullableUuidHolder(); + reader.read(targetHolder); + assertEquals(1, targetHolder.isSet); + assertEquals(0, targetHolder.start); + assertEquals(uuid, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start)); + } + } + + @Test + void testNullableUuidHolderReaderImplWithNull() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + vector.setNull(0); + vector.setValueCount(1); + + // Get null holder from vector + NullableUuidHolder sourceHolder = new NullableUuidHolder(); + vector.get(0, sourceHolder); + assertEquals(0, sourceHolder.isSet); + + // Create reader from null holder + NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder); + assertFalse(reader.isSet()); + assertNull(reader.readObject()); + + // Read into another holder + NullableUuidHolder targetHolder = new NullableUuidHolder(); + reader.read(targetHolder); + assertEquals(0, targetHolder.isSet); + } + } + + @Test + void testNullableUuidHolderReaderImplReadIntoUuidHolder() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid = UUID.randomUUID(); + vector.setSafe(0, uuid); + vector.setValueCount(1); + + // Get holder from vector + NullableUuidHolder sourceHolder = new NullableUuidHolder(); + vector.get(0, sourceHolder); + + // Create reader from holder + NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder); + + // Read into UuidHolder (non-nullable) + UuidHolder targetHolder = new UuidHolder(); + reader.read(targetHolder); + assertEquals(0, targetHolder.start); + assertEquals(uuid, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start)); + } + } + + @Test + void testNullableUuidHolderReaderImplWithNonZeroStart() throws Exception { + try (UuidVector vector = new UuidVector("test", allocator)) { + UUID uuid1 = UUID.randomUUID(); + UUID uuid2 = UUID.randomUUID(); + vector.setSafe(0, uuid1); + vector.setSafe(1, uuid2); + vector.setValueCount(2); + + // Get holder from index 1 (start = 16) + NullableUuidHolder sourceHolder = new NullableUuidHolder(); + vector.get(1, sourceHolder); + assertEquals(1, sourceHolder.isSet); + assertEquals(16, sourceHolder.start); + + // Create reader from holder + NullableUuidHolderReaderImpl reader = new NullableUuidHolderReaderImpl(sourceHolder); + assertEquals(uuid2, reader.readObject()); + + // Read into another holder and verify start is preserved + NullableUuidHolder targetHolder = new NullableUuidHolder(); + reader.read(targetHolder); + assertEquals(16, targetHolder.start); + assertEquals(uuid2, UuidUtility.uuidFromArrowBuf(targetHolder.buffer, targetHolder.start)); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java b/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java index daec331831..22c93b0cbe 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.apache.arrow.vector.TestUtils.newVarBinaryVector; import static org.apache.arrow.vector.TestUtils.newVarCharVector; import static org.apache.arrow.vector.TestUtils.newVector; @@ -56,6 +57,10 @@ import org.apache.arrow.vector.complex.impl.UnionListViewWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.holders.NullableIntHolder; +import org.apache.arrow.vector.holders.NullableTimeStampMicroTZHolder; +import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; +import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; +import org.apache.arrow.vector.holders.NullableTimeStampSecTZHolder; import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.holders.NullableVarBinaryHolder; import org.apache.arrow.vector.holders.NullableVarCharHolder; @@ -1233,7 +1238,7 @@ public void testSplitAndTransfer3() { // the size needed for the validity buffer final long validitySize = DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( - BaseValueVector.getValidityBufferSizeFromCount(2)); + getValidityBufferSizeFromCount(2)); assertEquals(allocatedMem + validitySize, allocator.getAllocatedMemory()); // The validity and offset buffers are sliced from a same buffer.See // BaseFixedWidthVector#allocateBytes. @@ -2464,7 +2469,7 @@ public void testDefaultAllocNewAll() { assertTrue(intVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * IntVector.TYPE_WIDTH) - + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity); + + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for BigIntVector. @@ -2473,7 +2478,7 @@ public void testDefaultAllocNewAll() { assertTrue(bigIntVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * bigIntVector.TYPE_WIDTH) - + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity); + + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for DecimalVector. @@ -2482,7 +2487,7 @@ public void testDefaultAllocNewAll() { assertTrue(decimalVector.getValueCapacity() >= defaultCapacity); expectedSize = (defaultCapacity * decimalVector.TYPE_WIDTH) - + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity); + + getValidityBufferSizeFromCount(defaultCapacity); assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); // verify that the wastage is within bounds for VarCharVector. @@ -2492,7 +2497,7 @@ public void testDefaultAllocNewAll() { assertTrue(varCharVector.getValueCapacity() >= defaultCapacity - 1); expectedSize = (defaultCapacity * VarCharVector.OFFSET_WIDTH) - + BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity) + + getValidityBufferSizeFromCount(defaultCapacity) + defaultCapacity * 8; // wastage should be less than 5%. assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); @@ -2501,7 +2506,7 @@ public void testDefaultAllocNewAll() { beforeSize = childAllocator.getAllocatedMemory(); bitVector.allocateNew(); assertTrue(bitVector.getValueCapacity() >= defaultCapacity); - expectedSize = BaseFixedWidthVector.getValidityBufferSizeFromCount(defaultCapacity) * 2; + expectedSize = getValidityBufferSizeFromCount(defaultCapacity) * 2; assertTrue(childAllocator.getAllocatedMemory() - beforeSize <= expectedSize * 1.05); } } @@ -2566,6 +2571,195 @@ public void testSetNullableVarCharHolderSafe() { } } + @Test + public void testTimeStampTZVectorSetSafeUnset() { + // reproduction of https://github.com/apache/arrow/issues/45084 + try (TimeStampMicroTZVector vector = new TimeStampMicroTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + // Set a valid value + NullableTimeStampMicroTZHolder validHolder = new NullableTimeStampMicroTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + // Unset the value using a holder with default (null) timezone + // The bug used to throw IllegalArgumentException because holder.timezone (null) != + // vector.timezone ("UTC") + // The correct behaviour is to not throw an exception and to unset the value. + NullableTimeStampMicroTZHolder unsetHolder = new NullableTimeStampMicroTZHolder(); + unsetHolder.isSet = 0; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampMilliTZVectorSetSafeUnset() { + // reproduction of https://github.com/apache/arrow/issues/45084 + try (TimeStampMilliTZVector vector = new TimeStampMilliTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampMilliTZHolder validHolder = new NullableTimeStampMilliTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampMilliTZHolder unsetHolder = new NullableTimeStampMilliTZHolder(); + unsetHolder.isSet = 0; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampNanoTZVectorSetSafeUnset() { + // reproduction of https://github.com/apache/arrow/issues/45084 + try (TimeStampNanoTZVector vector = new TimeStampNanoTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampNanoTZHolder validHolder = new NullableTimeStampNanoTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampNanoTZHolder unsetHolder = new NullableTimeStampNanoTZHolder(); + unsetHolder.isSet = 0; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampSecTZVectorSetSafeUnset() { + // reproduction of https://github.com/apache/arrow/issues/45084 + try (TimeStampSecTZVector vector = new TimeStampSecTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampSecTZHolder validHolder = new NullableTimeStampSecTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampSecTZHolder unsetHolder = new NullableTimeStampSecTZHolder(); + unsetHolder.isSet = 0; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampMicroTZVectorSetSafeUnsetExplicitTimezone() { + // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break + // workaround. + try (TimeStampMicroTZVector vector = new TimeStampMicroTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampMicroTZHolder validHolder = new NullableTimeStampMicroTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampMicroTZHolder unsetHolder = new NullableTimeStampMicroTZHolder(); + unsetHolder.isSet = 0; + unsetHolder.timezone = "UTC"; + + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampMilliTZVectorSetSafeUnsetExplicitTimezone() { + // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break + // workaround. + try (TimeStampMilliTZVector vector = new TimeStampMilliTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampMilliTZHolder validHolder = new NullableTimeStampMilliTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampMilliTZHolder unsetHolder = new NullableTimeStampMilliTZHolder(); + unsetHolder.isSet = 0; + unsetHolder.timezone = "UTC"; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampNanoTZVectorSetSafeUnsetExplicitTimezone() { + // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break + // workaround. + try (TimeStampNanoTZVector vector = new TimeStampNanoTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampNanoTZHolder validHolder = new NullableTimeStampNanoTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampNanoTZHolder unsetHolder = new NullableTimeStampNanoTZHolder(); + unsetHolder.isSet = 0; + unsetHolder.timezone = "UTC"; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + + @Test + public void testTimeStampSecTZVectorSetSafeUnsetExplicitTimezone() { + // Test to ensure fix added for https://github.com/apache/arrow/issues/45084 does not break + // workaround. + try (TimeStampSecTZVector vector = new TimeStampSecTZVector("vector", allocator, "UTC")) { + vector.allocateNew(); + + NullableTimeStampSecTZHolder validHolder = new NullableTimeStampSecTZHolder(); + validHolder.isSet = 1; + validHolder.value = 1000L; + validHolder.timezone = "UTC"; + vector.setSafe(0, validHolder); + + assertEquals(1000L, vector.get(0)); + + NullableTimeStampSecTZHolder unsetHolder = new NullableTimeStampSecTZHolder(); + unsetHolder.isSet = 0; + unsetHolder.timezone = "UTC"; + vector.setSafe(0, unsetHolder); + + assertNull(vector.getObject(0)); + } + } + @Test public void testSetNullableVarBinaryHolder() { try (VarBinaryVector vector = new VarBinaryVector("", allocator)) { @@ -3746,4 +3940,42 @@ public void testVectorLoadUnloadOnNonVariadicVectors() { } } } + + @Test + public void testEmptyVarCharOffsetBuffer() { + // Validates that offset buffer has at least OFFSET_WIDTH bytes (for offset[0]=0) + // even when valueCount is 0, per Arrow specification. + try (VarCharVector vector = newVarCharVector("varchar", allocator)) { + vector.allocateNew(); + vector.setValueCount(0); + + List buffers = vector.getFieldBuffers(); + // buffers: [validity, offset, data] + assertTrue( + buffers.get(1).readableBytes() >= BaseVariableWidthVector.OFFSET_WIDTH, + "Offset buffer should have at least " + + BaseVariableWidthVector.OFFSET_WIDTH + + " bytes for offset[0]"); + assertEquals(0, vector.getOffsetBuffer().getInt(0)); + } + } + + @Test + public void testEmptyLargeVarCharOffsetBuffer() { + // Validates that offset buffer has at least OFFSET_WIDTH bytes (for offset[0]=0) + // even when valueCount is 0, per Arrow specification. + try (LargeVarCharVector vector = new LargeVarCharVector("largevarchar", allocator)) { + vector.allocateNew(); + vector.setValueCount(0); + + List buffers = vector.getFieldBuffers(); + // buffers: [validity, offset, data] + assertTrue( + buffers.get(1).readableBytes() >= BaseLargeVariableWidthVector.OFFSET_WIDTH, + "Offset buffer should have at least " + + BaseLargeVariableWidthVector.OFFSET_WIDTH + + " bytes for offset[0]"); + assertEquals(0, vector.getOffsetBuffer().getLong(0)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java b/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java index 7a3a1bae63..baf5e672c8 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVariableWidthViewVector.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.apache.arrow.vector.TestUtils.newVector; import static org.apache.arrow.vector.TestUtils.newViewVarBinaryVector; import static org.apache.arrow.vector.TestUtils.newViewVarCharVector; @@ -60,6 +61,7 @@ import org.apache.arrow.vector.util.ReusableByteArray; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.validate.ValidateUtil; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -2367,7 +2369,7 @@ private void testSplitAndTransferOnValiditySplitHelper( // the allocation only consists in the size needed for the validity buffer final long validitySize = DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY.getRoundedSize( - BaseValueVector.getValidityBufferSizeFromCount(2)); + getValidityBufferSizeFromCount(2)); // we allocate view and data buffers for the target vector assertTrue(allocatedMem + validitySize < allocator.getAllocatedMemory()); // The validity is sliced from the same buffer.See BaseFixedWidthViewVector#allocateBytes. @@ -2444,7 +2446,7 @@ public void testSplitAndTransferWithLongStringsOnValiditySplit() { final ViewVarBinaryVector sourceVector = newViewVarBinaryVector(EMPTY_SCHEMA_PATH, allocator)) { testSplitAndTransferOnValiditySplitHelper( - targetVector, sourceVector, startIndex, length, data); + targetVector, sourceVector, startIndex, length, binaryData); } } @@ -2851,4 +2853,18 @@ public void testVectorLoadUnloadOnMixedTypes() { } } } + + @Test + public void testValidate() { + try (final ViewVarCharVector vector = new ViewVarCharVector("v", allocator)) { + vector.validateFull(); + setVector(vector, STR1, STR2, STR3); + vector.validateFull(); + + vector.getDataBuffer().capacity(0); + ValidateUtil.ValidateException e = + assertThrows(ValidateUtil.ValidateException.class, () -> vector.validate()); + assertTrue(e.getMessage().contains("Not enough capacity for data buffer")); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java index c121d94892..bd3113f8bc 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java @@ -171,6 +171,26 @@ public void testAddVector() { } } + @Test + public void testAddVectorAtEnd() { + try (final IntVector intVector1 = new IntVector("intVector1", allocator); + final IntVector intVector2 = new IntVector("intVector2", allocator); + final IntVector intVector3 = new IntVector("intVector3", allocator); ) { + + VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2)); + assertEquals(2, original.getFieldVectors().size()); + + VectorSchemaRoot newRecordBatch = original.addVector(2, intVector3); + assertEquals(3, newRecordBatch.getFieldVectors().size()); + assertEquals(intVector1, newRecordBatch.getFieldVectors().get(0)); + assertEquals(intVector2, newRecordBatch.getFieldVectors().get(1)); + assertEquals(intVector3, newRecordBatch.getFieldVectors().get(2)); + + original.close(); + newRecordBatch.close(); + } + } + @Test public void testRemoveVector() { try (final IntVector intVector1 = new IntVector("intVector1", allocator); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java index 6121fb67fe..782535fccc 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloadLoad.java @@ -17,6 +17,7 @@ package org.apache.arrow.vector; import static java.util.Arrays.asList; +import static org.apache.arrow.vector.BitVectorHelper.getValidityBufferSizeFromCount; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -215,7 +216,7 @@ public void testLoadValidityBuffer() throws IOException { int count = 10; ArrowBuf[] values = new ArrowBuf[4]; for (int i = 0; i < 4; i += 2) { - ArrowBuf buf1 = allocator.buffer(BitVectorHelper.getValidityBufferSize(count)); + ArrowBuf buf1 = allocator.buffer(getValidityBufferSizeFromCount(count)); ArrowBuf buf2 = allocator.buffer(count * 4); // integers buf1.setZero(0, buf1.capacity()); buf2.setZero(0, buf2.capacity()); diff --git a/vector/src/test/java/org/apache/arrow/vector/UuidVector.java b/vector/src/test/java/org/apache/arrow/vector/UuidVector.java deleted file mode 100644 index 5c90d45f60..0000000000 --- a/vector/src/test/java/org/apache/arrow/vector/UuidVector.java +++ /dev/null @@ -1,114 +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.arrow.vector; - -import java.nio.ByteBuffer; -import java.util.UUID; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.util.hash.ArrowBufHasher; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.UuidType; -import org.apache.arrow.vector.util.TransferPair; - -public class UuidVector extends ExtensionTypeVector - implements ValueIterableVector { - private final Field field; - - public UuidVector( - String name, BufferAllocator allocator, FixedSizeBinaryVector underlyingVector) { - super(name, allocator, underlyingVector); - this.field = new Field(name, FieldType.nullable(new UuidType()), null); - } - - public UuidVector(String name, BufferAllocator allocator) { - super(name, allocator, new FixedSizeBinaryVector(name, allocator, 16)); - this.field = new Field(name, FieldType.nullable(new UuidType()), null); - } - - @Override - public UUID getObject(int index) { - final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index)); - return new UUID(bb.getLong(), bb.getLong()); - } - - @Override - public int hashCode(int index) { - return hashCode(index, null); - } - - @Override - public int hashCode(int index, ArrowBufHasher hasher) { - return getUnderlyingVector().hashCode(index, hasher); - } - - public void set(int index, UUID uuid) { - ByteBuffer bb = ByteBuffer.allocate(16); - bb.putLong(uuid.getMostSignificantBits()); - bb.putLong(uuid.getLeastSignificantBits()); - getUnderlyingVector().set(index, bb.array()); - } - - @Override - public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) { - getUnderlyingVector() - .copyFromSafe(fromIndex, thisIndex, ((UuidVector) from).getUnderlyingVector()); - } - - @Override - public Field getField() { - return field; - } - - @Override - public TransferPair makeTransferPair(ValueVector to) { - return new TransferImpl((UuidVector) to); - } - - public void setSafe(int index, byte[] value) { - getUnderlyingVector().setIndexDefined(index); - getUnderlyingVector().setSafe(index, value); - } - - public class TransferImpl implements TransferPair { - UuidVector to; - ValueVector targetUnderlyingVector; - TransferPair tp; - - public TransferImpl(UuidVector to) { - this.to = to; - targetUnderlyingVector = this.to.getUnderlyingVector(); - tp = getUnderlyingVector().makeTransferPair(targetUnderlyingVector); - } - - public UuidVector getTo() { - return this.to; - } - - public void transfer() { - tp.transfer(); - } - - public void splitAndTransfer(int startIndex, int length) { - tp.splitAndTransfer(startIndex, length); - } - - public void copyValueSafe(int fromIndex, int toIndex) { - tp.copyValueSafe(fromIndex, toIndex); - } - } -} diff --git a/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java b/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java index 08da786eb2..9624734356 100644 --- a/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java +++ b/vector/src/test/java/org/apache/arrow/vector/compare/TestRangeEqualsVisitor.java @@ -22,6 +22,7 @@ import java.nio.charset.Charset; import java.util.Arrays; +import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; @@ -39,6 +40,7 @@ import org.apache.arrow.vector.complex.LargeListViewVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.ListViewVector; +import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.impl.NullableStructWriter; @@ -53,7 +55,9 @@ import org.apache.arrow.vector.holders.NullableUInt4Holder; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.Types; +import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.ArrowType.RunEndEncoded; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.junit.jupiter.api.AfterEach; @@ -1003,6 +1007,54 @@ public void testLargeListViewVectorApproxEquals() { } } + @Test + public void testRunEndEncodedFloat8ApproxEquals() { + try (final Float8Vector vector1 = new Float8Vector("float", allocator); + final Float8Vector vector2 = new Float8Vector("float", allocator); + final Float8Vector vector3 = new Float8Vector("float", allocator); + final IntVector reeVector = new IntVector("ree", allocator)) { + + final float epsilon = 1.0E-6f; + setVector(vector1, 1.1, 2.2); + setVector(vector2, 1.1 + epsilon / 2, 2.2 + epsilon / 2); + setVector(vector3, 1.1 + epsilon * 2, 2.2 + epsilon * 2); + setVector(reeVector, 1, 3); + + ArrowType type = MinorType.FLOAT8.getType(); + final FieldType valueType = FieldType.notNullable(type); + final FieldType runEndType = FieldType.notNullable(MinorType.INT.getType()); + + final Field valueField = new Field("value", valueType, null); + final Field runEndField = new Field("ree", runEndType, null); + + Field field = + new Field( + "ree_float", + FieldType.notNullable(RunEndEncoded.INSTANCE), + List.of(runEndField, valueField)); + + try (final RunEndEncodedVector encodedVector1 = + new RunEndEncodedVector(field, allocator, reeVector, vector1, null); + final RunEndEncodedVector encodedVector2 = + new RunEndEncodedVector(field, allocator, reeVector, vector2, null); + final RunEndEncodedVector encodedVector3 = + new RunEndEncodedVector(field, allocator, reeVector, vector3, null)) { + + encodedVector1.setValueCount(3); + encodedVector2.setValueCount(3); + encodedVector3.setValueCount(3); + + Range range = new Range(0, 0, encodedVector1.getValueCount()); + assertTrue( + new ApproxEqualsVisitor(encodedVector1, encodedVector2, epsilon, epsilon) + .rangeEquals(range)); + assertFalse( + new ApproxEqualsVisitor(encodedVector1, encodedVector3, epsilon, epsilon) + .rangeEquals(range)); + } + } + } + private void writeStructVector(NullableStructWriter writer, int value1, long value2) { writer.start(); writer.integer("f0").writeInt(value1); diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java index 3bc02c6029..b2a8cf9ba4 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java @@ -20,6 +20,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; +import java.util.UUID; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.DecimalVector; @@ -30,8 +31,10 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.complex.writer.BaseWriter; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.DecimalHolder; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -845,4 +848,110 @@ public void testCopyMapVectorWithMapValue() { assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); } } + + @Test + public void testCopyListVectorWithExtensionType() { + try (ListVector from = ListVector.empty("v", allocator); + ListVector to = ListVector.empty("v", allocator)) { + + UnionListWriter listWriter = from.getWriter(); + listWriter.allocate(); + + for (int i = 0; i < COUNT; i++) { + listWriter.setPosition(i); + listWriter.startList(); + ExtensionWriter extensionWriter = listWriter.extension(UuidType.INSTANCE); + extensionWriter.writeExtension(UUID.randomUUID()); + extensionWriter.writeExtension(UUID.randomUUID()); + listWriter.endList(); + } + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + + to.setValueCount(COUNT); + + // validate equals + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } + + @Test + public void testCopyMapVectorWithExtensionType() { + try (final MapVector from = MapVector.empty("v", allocator, false); + final MapVector to = MapVector.empty("v", allocator, false)) { + + from.allocateNew(); + + UnionMapWriter mapWriter = from.getWriter(); + for (int i = 0; i < COUNT; i++) { + mapWriter.setPosition(i); + mapWriter.startMap(); + mapWriter.startEntry(); + ExtensionWriter extensionKeyWriter = mapWriter.key().extension(UuidType.INSTANCE); + extensionKeyWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE); + ExtensionWriter extensionValueWriter = mapWriter.value().extension(UuidType.INSTANCE); + extensionValueWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE); + mapWriter.endEntry(); + mapWriter.endMap(); + } + + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + to.setValueCount(COUNT); + + // validate equals + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } + + @Test + public void testCopyStructVectorWithExtensionType() { + try (final StructVector from = StructVector.empty("v", allocator); + final StructVector to = StructVector.empty("v", allocator)) { + + from.allocateNewSafe(); + + NullableStructWriter structWriter = from.getWriter(); + for (int i = 0; i < COUNT; i++) { + structWriter.setPosition(i); + structWriter.start(); + ExtensionWriter extensionWriter1 = structWriter.extension("uuid1", UuidType.INSTANCE); + extensionWriter1.writeExtension(UUID.randomUUID(), UuidType.INSTANCE); + ExtensionWriter extensionWriter2 = structWriter.extension("uuid2", UuidType.INSTANCE); + extensionWriter2.writeExtension(UUID.randomUUID(), UuidType.INSTANCE); + structWriter.end(); + } + + from.setValueCount(COUNT); + + // copy values + FieldReader in = from.getReader(); + FieldWriter out = to.getWriter(); + for (int i = 0; i < COUNT; i++) { + in.setPosition(i); + out.setPosition(i); + ComplexCopier.copy(in, out); + } + to.setValueCount(COUNT); + + // validate equals + assertTrue(VectorEqualsVisitor.vectorEquals(from, to)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java index 1556852c5a..5b6d65d6ba 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java @@ -31,6 +31,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DirtyRootAllocator; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; import org.apache.arrow.vector.UuidVector; @@ -41,6 +42,7 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; import org.apache.arrow.vector.holders.NullableDecimalHolder; @@ -48,15 +50,16 @@ import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; import org.apache.arrow.vector.holders.UnionHolder; +import org.apache.arrow.vector.holders.UuidHolder; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.UuidType; import org.apache.arrow.vector.util.DecimalUtility; import org.apache.arrow.vector.util.Text; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -100,7 +103,6 @@ public void testPromoteToUnion() throws Exception { writer.integer("A").writeInt(10); // we don't write anything in 3 - writer.setPosition(4); writer.integer("A").writeInt(100); @@ -130,9 +132,21 @@ public void testPromoteToUnion() throws Exception { binHolder.buffer = buf; writer.fixedSizeBinary("A", 4).write(binHolder); + writer.setPosition(9); + UUID uuid = UUID.randomUUID(); + writer.extension("A", UuidType.INSTANCE).writeExtension(uuid, UuidType.INSTANCE); + writer.end(); + + writer.setPosition(10); + UUID uuid2 = UUID.randomUUID(); + UuidHolder uuidHolder = new UuidHolder(); + uuidHolder.buffer = allocator.buffer(UuidType.UUID_BYTE_WIDTH); + uuidHolder.buffer.setBytes(0, UuidUtility.getBytesFromUUID(uuid2)); + writer.extension("A", UuidType.INSTANCE).write(uuidHolder); writer.end(); + allocator.releaseBytes(UuidType.UUID_BYTE_WIDTH); - container.setValueCount(9); + container.setValueCount(11); final UnionVector uv = v.getChild("A", UnionVector.class); @@ -169,6 +183,12 @@ public void testPromoteToUnion() throws Exception { .order(ByteOrder.nativeOrder()) .getInt()); + assertFalse(uv.isNull(9), "9 shouldn't be null"); + assertEquals(uuid, uv.getObject(9)); + + assertFalse(uv.isNull(10), "10 shouldn't be null"); + assertEquals(uuid2, uv.getObject(10)); + container.clear(); container.allocateNew(); @@ -785,18 +805,17 @@ public void testExtensionType() throws Exception { try (final NonNullableStructVector container = NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); final UuidVector v = - container.addOrGet("uuid", FieldType.nullable(new UuidType()), UuidVector.class); + container.addOrGet("uuid", FieldType.nullable(UuidType.INSTANCE), UuidVector.class); final PromotableWriter writer = new PromotableWriter(v, container)) { UUID u1 = UUID.randomUUID(); UUID u2 = UUID.randomUUID(); container.allocateNew(); container.setValueCount(1); - writer.addExtensionTypeWriterFactory(new UuidWriterFactory()); writer.setPosition(0); - writer.writeExtension(u1); + writer.writeExtension(u1, UuidType.INSTANCE); writer.setPosition(1); - writer.writeExtension(u2); + writer.writeExtension(u2, UuidType.INSTANCE); container.setValueCount(2); @@ -805,4 +824,29 @@ public void testExtensionType() throws Exception { assertEquals(u2, uuidVector.getObject(1)); } } + + @Test + public void testExtensionTypeForList() throws Exception { + try (final ListVector container = ListVector.empty(EMPTY_SCHEMA_PATH, allocator); + final UuidVector v = + (UuidVector) + container.addOrGetVector(FieldType.nullable(UuidType.INSTANCE)).getVector(); + final PromotableWriter writer = new PromotableWriter(v, container)) { + UUID u1 = UUID.randomUUID(); + UUID u2 = UUID.randomUUID(); + container.allocateNew(); + container.setValueCount(1); + + writer.setPosition(0); + writer.writeExtension(u1, UuidType.INSTANCE); + writer.setPosition(1); + writer.writeExtension(u2, UuidType.INSTANCE); + + container.setValueCount(2); + + FieldVector uuidVector = container.getDataVector(); + assertEquals(u1, uuidVector.getObject(0)); + assertEquals(u2, uuidVector.getObject(1)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java index 2745386db4..80d03cae6d 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -31,6 +32,7 @@ import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -71,9 +73,11 @@ import org.apache.arrow.vector.complex.reader.Float8Reader; import org.apache.arrow.vector.complex.reader.IntReader; import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter; +import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter; import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.holders.DecimalHolder; import org.apache.arrow.vector.holders.DurationHolder; import org.apache.arrow.vector.holders.FixedSizeBinaryHolder; @@ -82,8 +86,11 @@ import org.apache.arrow.vector.holders.NullableFixedSizeBinaryHolder; import org.apache.arrow.vector.holders.NullableTimeStampMilliTZHolder; import org.apache.arrow.vector.holders.NullableTimeStampNanoTZHolder; +import org.apache.arrow.vector.holders.NullableUuidHolder; import org.apache.arrow.vector.holders.TimeStampMilliTZHolder; +import org.apache.arrow.vector.holders.UuidHolder; import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.ArrowType.ArrowTypeID; import org.apache.arrow.vector.types.pojo.ArrowType.Int; @@ -99,6 +106,7 @@ import org.apache.arrow.vector.util.JsonStringHashMap; import org.apache.arrow.vector.util.Text; import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.util.UuidUtility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -1098,6 +1106,13 @@ public void simpleUnion() throws Exception { new UnionVector("union", allocator, /* field type */ null, /* call-back */ null); UnionWriter unionWriter = new UnionWriter(vector); unionWriter.allocate(); + + UUID uuid = UUID.randomUUID(); + ByteBuffer bb = ByteBuffer.allocate(16); + bb.putLong(uuid.getMostSignificantBits()); + bb.putLong(uuid.getLeastSignificantBits()); + byte[] uuidByte = bb.array(); + for (int i = 0; i < COUNT; i++) { unionWriter.setPosition(i); if (i % 5 == 0) { @@ -1120,6 +1135,12 @@ public void simpleUnion() throws Exception { holder.buffer = buf; unionWriter.write(holder); bufs.add(buf); + } else if (i % 5 == 4) { + UuidHolder holder = new UuidHolder(); + holder.buffer = allocator.buffer(UuidType.UUID_BYTE_WIDTH); + holder.buffer.setBytes(0, uuidByte); + unionWriter.write(holder); + allocator.releaseBytes(UuidType.UUID_BYTE_WIDTH); } else { unionWriter.writeFloat4((float) i); } @@ -1145,6 +1166,10 @@ public void simpleUnion() throws Exception { unionReader.read(holder); assertEquals(i, holder.buffer.getInt(0)); assertEquals(4, holder.byteWidth); + } else if (i % 5 == 4) { + NullableUuidHolder holder = new NullableUuidHolder(); + unionReader.read(holder); + assertEquals(UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start), uuid); } else { assertEquals((float) i, unionReader.readFloat(), 1e-12); } @@ -2489,4 +2514,79 @@ public void unionWithVarCharAndBinaryHelpers() throws Exception { "row12", new String(vector.getLargeVarBinaryVector().get(11), StandardCharsets.UTF_8)); } } + + @Test + public void extensionWriterReader() throws Exception { + // test values + UUID u1 = UUID.randomUUID(); + + try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) { + // write + + ComplexWriter writer = new ComplexWriterImpl("root", parent); + StructWriter rootWriter = writer.rootAsStruct(); + + { + ExtensionWriter extensionWriter = rootWriter.extension("uuid1", UuidType.INSTANCE); + extensionWriter.setPosition(0); + extensionWriter.writeExtension(u1, UuidType.INSTANCE); + } + // read + StructReader rootReader = new SingleStructReaderImpl(parent).reader("root"); + { + FieldReader uuidReader = rootReader.reader("uuid1"); + uuidReader.setPosition(0); + NullableUuidHolder uuidHolder = new NullableUuidHolder(); + uuidReader.read(uuidHolder); + UUID actualUuid = UuidUtility.uuidFromArrowBuf(uuidHolder.buffer, 0); + assertEquals(u1, actualUuid); + assertTrue(uuidReader.isSet()); + assertEquals(uuidReader.getMinorType(), MinorType.EXTENSIONTYPE); + assertInstanceOf(UuidType.class, uuidReader.getField().getFieldType().getType()); + } + } + } + + @Test + void testListOfDenseUnionWriterNPE() { + // Regression test for https://github.com/apache/arrow-java/issues/399 + try (ListVector listVector = ListVector.empty("list", allocator)) { + listVector.addOrGetVector(FieldType.nullable(MinorType.DENSEUNION.getType())); + UnionListWriter listWriter = listVector.getWriter(); + + listWriter.startList(); + listWriter.endList(); + } + } + + @Test + void testListOfDenseUnionWriterWithData() { + try (ListVector listVector = ListVector.empty("list", allocator)) { + listVector.addOrGetVector(FieldType.nullable(MinorType.DENSEUNION.getType())); + + UnionListWriter listWriter = listVector.getWriter(); + listWriter.startList(); + listWriter.writeInt(100); + listWriter.writeBigInt(200L); + listWriter.endList(); + + listWriter.startList(); + listWriter.writeFloat4(3.14f); + listWriter.endList(); + + listVector.setValueCount(2); + + assertEquals(2, listVector.getValueCount()); + + List value0 = (List) listVector.getObject(0); + List value1 = (List) listVector.getObject(1); + + assertEquals(2, value0.size()); + assertEquals(100, value0.get(0)); + assertEquals(200L, value0.get(1)); + + assertEquals(1, value1.size()); + assertEquals(3.14f, value1.get(0)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java index bf1b9b0dfa..5bb5962704 100644 --- a/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java +++ b/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestSimpleWriter.java @@ -20,20 +20,16 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import java.nio.ByteBuffer; -import java.util.UUID; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.LargeVarBinaryVector; import org.apache.arrow.vector.LargeVarCharVector; -import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.impl.LargeVarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.LargeVarCharWriterImpl; -import org.apache.arrow.vector.complex.impl.UuidWriterImpl; import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.VarCharWriterImpl; -import org.apache.arrow.vector.holder.UuidHolder; import org.apache.arrow.vector.util.Text; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -188,20 +184,4 @@ public void testWriteTextToLargeVarChar() throws Exception { assertEquals(input, result); } } - - @Test - public void testWriteToExtensionVector() throws Exception { - try (UuidVector vector = new UuidVector("test", allocator); - UuidWriterImpl writer = new UuidWriterImpl(vector)) { - UUID uuid = UUID.randomUUID(); - ByteBuffer bb = ByteBuffer.allocate(16); - bb.putLong(uuid.getMostSignificantBits()); - bb.putLong(uuid.getLeastSignificantBits()); - UuidHolder holder = new UuidHolder(); - holder.value = bb.array(); - writer.write(holder); - UUID result = vector.getObject(0); - assertEquals(uuid, result); - } - } } diff --git a/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java b/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java index d24708d66c..ae5ac0726c 100644 --- a/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java +++ b/vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector.types.pojo; +import static org.apache.arrow.vector.TestUtils.ensureRegistered; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -43,10 +44,13 @@ import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.UuidVector; import org.apache.arrow.vector.ValueIterableVector; +import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.complex.writer.FieldWriter; +import org.apache.arrow.vector.extension.UuidType; import org.apache.arrow.vector.ipc.ArrowFileReader; import org.apache.arrow.vector.ipc.ArrowFileWriter; import org.apache.arrow.vector.types.FloatingPointPrecision; @@ -59,9 +63,9 @@ public class TestExtensionType { /** Test that a custom UUID type can be round-tripped through a temporary file. */ @Test public void roundtripUuid() throws IOException { - ExtensionTypeRegistry.register(new UuidType()); + ensureRegistered(UuidType.INSTANCE); final Schema schema = - new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); + new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { UUID u1 = UUID.randomUUID(); @@ -89,7 +93,7 @@ public void roundtripUuid() throws IOException { assertEquals(root.getSchema(), readerRoot.getSchema()); final Field field = readerRoot.getSchema().getFields().get(0); - final UuidType expectedType = new UuidType(); + final UuidType expectedType = UuidType.INSTANCE; assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), expectedType.extensionName()); @@ -113,9 +117,9 @@ public void roundtripUuid() throws IOException { /** Test that a custom UUID type can be read as its underlying type. */ @Test public void readUnderlyingType() throws IOException { - ExtensionTypeRegistry.register(new UuidType()); + ensureRegistered(UuidType.INSTANCE); final Schema schema = - new Schema(Collections.singletonList(Field.nullable("a", new UuidType()))); + new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE))); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { UUID u1 = UUID.randomUUID(); @@ -135,7 +139,7 @@ public void readUnderlyingType() throws IOException { writer.end(); } - ExtensionTypeRegistry.unregister(new UuidType()); + ExtensionTypeRegistry.unregister(UuidType.INSTANCE); try (final SeekableByteChannel channel = Files.newByteChannel(Paths.get(file.getAbsolutePath())); @@ -153,7 +157,7 @@ public void readUnderlyingType() throws IOException { .getByteWidth()); final Field field = readerRoot.getSchema().getFields().get(0); - final UuidType expectedType = new UuidType(); + final UuidType expectedType = UuidType.INSTANCE; assertEquals( field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME), expectedType.extensionName()); @@ -254,7 +258,7 @@ public void roundtripLocation() throws IOException { @Test public void testVectorCompare() { - UuidType uuidType = new UuidType(); + UuidType uuidType = UuidType.INSTANCE; ExtensionTypeRegistry.register(uuidType); try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE); UuidVector a1 = @@ -331,6 +335,11 @@ public String serialize() { public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { return new LocationVector(name, allocator); } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + throw new UnsupportedOperationException("Not yet implemented."); + } } public static class LocationVector extends ExtensionTypeVector diff --git a/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java b/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java deleted file mode 100644 index 5e2bd8881b..0000000000 --- a/vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java +++ /dev/null @@ -1,60 +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.arrow.vector.types.pojo; - -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.FixedSizeBinaryVector; -import org.apache.arrow.vector.UuidVector; -import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType; - -public class UuidType extends ExtensionType { - - @Override - public ArrowType storageType() { - return new ArrowType.FixedSizeBinary(16); - } - - @Override - public String extensionName() { - return "uuid"; - } - - @Override - public boolean extensionEquals(ExtensionType other) { - return other instanceof UuidType; - } - - @Override - public ArrowType deserialize(ArrowType storageType, String serializedData) { - if (!storageType.equals(storageType())) { - throw new UnsupportedOperationException( - "Cannot construct UuidType from underlying type " + storageType); - } - return new UuidType(); - } - - @Override - public String serialize() { - return ""; - } - - @Override - public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) { - return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16)); - } -} diff --git a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java index 4ee9630a4d..9a8143f51b 100644 --- a/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java +++ b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java @@ -26,10 +26,13 @@ import java.util.List; import java.util.stream.IntStream; import java.util.stream.Stream; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseValueVector; +import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; @@ -47,11 +50,13 @@ import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.LargeListVector; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.RunEndEncodedVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.complex.UnionVector; import org.apache.arrow.vector.holders.NullableBigIntHolder; import org.apache.arrow.vector.holders.NullableFloat4Holder; import org.apache.arrow.vector.holders.NullableIntHolder; +import org.apache.arrow.vector.ipc.message.ArrowFieldNode; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -177,6 +182,82 @@ public void testAppendVariableWidthVector() { } } + @Test + public void testAppendVariableWidthVectorWithNonZeroStartOffset() { + try (VarCharVector target = new VarCharVector("", allocator); + VarCharVector delta = new VarCharVector("", allocator)) { + + target.allocateNew(64, 4); + ValueVectorDataPopulator.setVector(target, "a0", "a1"); + + // Build a delta vector whose offset buffer does not start at zero, as produced e.g. by + // importing a sliced array through the C data interface. The values are "BBBB" and + // "CCCC"; the data buffer additionally holds 4 bytes of unreferenced prefix ("AAAA"). + try (ArrowBuf validity = allocator.buffer(1); + ArrowBuf offsets = allocator.buffer(12); + ArrowBuf data = allocator.buffer(12)) { + validity.setByte(0, 0b11); + offsets.setInt(0, 4); + offsets.setInt(4, 8); + offsets.setInt(8, 12); + data.setBytes(0, "AAAABBBBCCCC".getBytes(StandardCharsets.UTF_8)); + delta.loadFieldBuffers(new ArrowFieldNode(2, 0), Arrays.asList(validity, offsets, data)); + } + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + // the unreferenced prefix must not be appended + assertEquals( + 4 + 8, + target + .getOffsetBuffer() + .getInt((long) target.getValueCount() * BaseVariableWidthVector.OFFSET_WIDTH)); + + try (VarCharVector expected = new VarCharVector("expected", allocator)) { + expected.allocateNew(); + ValueVectorDataPopulator.setVector(expected, "a0", "a1", "BBBB", "CCCC"); + assertVectorsEqual(expected, target); + } + } + } + + @Test + public void testAppendLargeVariableWidthVectorWithNonZeroStartOffset() { + try (LargeVarCharVector target = new LargeVarCharVector("", allocator); + LargeVarCharVector delta = new LargeVarCharVector("", allocator)) { + + target.allocateNew(64, 4); + ValueVectorDataPopulator.setVector(target, "a0", "a1"); + + try (ArrowBuf validity = allocator.buffer(1); + ArrowBuf offsets = allocator.buffer(24); + ArrowBuf data = allocator.buffer(12)) { + validity.setByte(0, 0b11); + offsets.setLong(0, 4); + offsets.setLong(8, 8); + offsets.setLong(16, 12); + data.setBytes(0, "AAAABBBBCCCC".getBytes(StandardCharsets.UTF_8)); + delta.loadFieldBuffers(new ArrowFieldNode(2, 0), Arrays.asList(validity, offsets, data)); + } + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals( + 4 + 8, + target + .getOffsetBuffer() + .getLong((long) target.getValueCount() * BaseLargeVariableWidthVector.OFFSET_WIDTH)); + + try (LargeVarCharVector expected = new LargeVarCharVector("expected", allocator)) { + expected.allocateNew(); + ValueVectorDataPopulator.setVector(expected, "a0", "a1", "BBBB", "CCCC"); + assertVectorsEqual(expected, target); + } + } + } + @Test public void testAppendVariableWidthViewVector() { final int length1 = 10; @@ -430,6 +511,115 @@ public void testAppendListVector() { } } + @Test + public void testAppendListVectorWithNonZeroStartOffset() { + try (ListVector target = ListVector.empty("target", allocator); + ListVector delta = ListVector.empty("delta", allocator)) { + + target.allocateNew(); + ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1), Arrays.asList(2, 3)); + + // Build a delta vector whose offset buffer does not start at zero, as produced e.g. by + // importing a sliced array through the C data interface: lists [10, 11] and [12, 13], + // with one unreferenced prefix element (9) in the data vector. + delta.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); + IntVector deltaDataVector = (IntVector) delta.getDataVector(); + deltaDataVector.allocateNew(5); + for (int i = 0; i < 5; i++) { + deltaDataVector.set(i, 9 + i); + } + deltaDataVector.setValueCount(5); + try (ArrowBuf validity = allocator.buffer(1); + ArrowBuf offsets = allocator.buffer(12)) { + validity.setByte(0, 0b11); + offsets.setInt(0, 1); + offsets.setInt(4, 3); + offsets.setInt(8, 5); + delta.loadFieldBuffers(new ArrowFieldNode(2, 0), Arrays.asList(validity, offsets)); + } + assertEquals(Arrays.asList(10, 11), delta.getObject(0)); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(4, target.getValueCount()); + // the unreferenced prefix element must not be appended + assertEquals( + 4 + 4, + target.getOffsetBuffer().getInt((long) target.getValueCount() * ListVector.OFFSET_WIDTH)); + assertEquals(Arrays.asList(0, 1), target.getObject(0)); + assertEquals(Arrays.asList(2, 3), target.getObject(1)); + assertEquals(Arrays.asList(10, 11), target.getObject(2)); + assertEquals(Arrays.asList(12, 13), target.getObject(3)); + } + } + + @Test + public void testAppendLargeListVector() { + try (LargeListVector target = LargeListVector.empty("target", allocator); + LargeListVector delta = LargeListVector.empty("delta", allocator)) { + + target.allocateNew(); + ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1), null, Arrays.asList(4, 5)); + + delta.allocateNew(); + ValueVectorDataPopulator.setVector(delta, Arrays.asList(10, 11, 12), Arrays.asList(13, 14)); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(5, target.getValueCount()); + assertEquals(Arrays.asList(0, 1), target.getObject(0)); + assertTrue(target.isNull(1)); + assertEquals(Arrays.asList(4, 5), target.getObject(2)); + assertEquals(Arrays.asList(10, 11, 12), target.getObject(3)); + assertEquals(Arrays.asList(13, 14), target.getObject(4)); + } + } + + @Test + public void testAppendLargeListVectorWithNonZeroStartOffset() { + try (LargeListVector target = LargeListVector.empty("target", allocator); + LargeListVector delta = LargeListVector.empty("delta", allocator)) { + + target.allocateNew(); + ValueVectorDataPopulator.setVector(target, Arrays.asList(0, 1), Arrays.asList(2, 3)); + + // same as testAppendListVectorWithNonZeroStartOffset, with 8-byte offsets + delta.addOrGetVector(FieldType.nullable(Types.MinorType.INT.getType())); + IntVector deltaDataVector = (IntVector) delta.getDataVector(); + deltaDataVector.allocateNew(5); + for (int i = 0; i < 5; i++) { + deltaDataVector.set(i, 9 + i); + } + deltaDataVector.setValueCount(5); + try (ArrowBuf validity = allocator.buffer(1); + ArrowBuf offsets = allocator.buffer(24)) { + validity.setByte(0, 0b11); + offsets.setLong(0, 1); + offsets.setLong(8, 3); + offsets.setLong(16, 5); + delta.loadFieldBuffers(new ArrowFieldNode(2, 0), Arrays.asList(validity, offsets)); + } + assertEquals(Arrays.asList(10, 11), delta.getObject(0)); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(4, target.getValueCount()); + // the unreferenced prefix element must not be appended + assertEquals( + 4 + 4, + target + .getOffsetBuffer() + .getLong((long) target.getValueCount() * LargeListVector.OFFSET_WIDTH)); + assertEquals(Arrays.asList(0, 1), target.getObject(0)); + assertEquals(Arrays.asList(2, 3), target.getObject(1)); + assertEquals(Arrays.asList(10, 11), target.getObject(2)); + assertEquals(Arrays.asList(12, 13), target.getObject(3)); + } + } + @Test public void testAppendEmptyListVector() { try (ListVector target = ListVector.empty("target", allocator); @@ -1025,6 +1215,72 @@ public void testAppendDenseUnionVectorMismatch() { } } + @Test + public void testAppendRunEndEncodedVector() { + final FieldType reeFieldType = FieldType.notNullable(ArrowType.RunEndEncoded.INSTANCE); + final Field runEndsField = + new Field("runEnds", FieldType.notNullable(Types.MinorType.INT.getType()), null); + final Field valuesField = Field.nullable("values", Types.MinorType.INT.getType()); + final List children = Arrays.asList(runEndsField, valuesField); + + final Field targetField = new Field("target", reeFieldType, children); + final Field deltaField = new Field("delta", reeFieldType, children); + try (RunEndEncodedVector target = new RunEndEncodedVector(targetField, allocator, null); + RunEndEncodedVector delta = new RunEndEncodedVector(deltaField, allocator, null)) { + + // populate target + target.allocateNew(); + // data: [1, 1, 2, null, 3, 3, 3] (7 values) + // values: [1, 2, null, 3] + // runEnds: [2, 3, 4, 7] + ValueVectorDataPopulator.setVector((IntVector) target.getValuesVector(), 1, 2, null, 3); + ValueVectorDataPopulator.setVector((IntVector) target.getRunEndsVector(), 2, 3, 4, 7); + target.setValueCount(7); + + // populate delta + delta.allocateNew(); + // data: [3, 4, 4, 5, null, null] (6 values) + // values: [3, 4, 5, null] + // runEnds: [1, 3, 4, 6] + ValueVectorDataPopulator.setVector((IntVector) delta.getValuesVector(), 3, 4, 5, null); + ValueVectorDataPopulator.setVector((IntVector) delta.getRunEndsVector(), 1, 3, 4, 6); + delta.setValueCount(6); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(13, target.getValueCount()); + + final Field expectedField = new Field("expected", reeFieldType, children); + try (RunEndEncodedVector expected = new RunEndEncodedVector(expectedField, allocator, null)) { + expected.allocateNew(); + // expected data: [1, 1, 2, null, 3, 3, 3, 3, 4, 4, 5, null, null] (13 values) + // expected values: [1, 2, null, 3, 3, 4, 5, null] + // expected runEnds: [2, 3, 4, 7, 8, 10, 11, 13] + ValueVectorDataPopulator.setVector( + (IntVector) expected.getValuesVector(), 1, 2, null, 3, 3, 4, 5, null); + ValueVectorDataPopulator.setVector( + (IntVector) expected.getRunEndsVector(), 2, 3, 4, 7, 8, 10, 11, 13); + expected.setValueCount(13); + + assertVectorsEqual(expected, target); + } + + // Check that delta is unchanged. + final Field expectedDeltaField = new Field("expectedDelta", reeFieldType, children); + try (RunEndEncodedVector expectedDelta = + new RunEndEncodedVector(expectedDeltaField, allocator, null)) { + expectedDelta.allocateNew(); + ValueVectorDataPopulator.setVector( + (IntVector) expectedDelta.getValuesVector(), 3, 4, 5, null); + ValueVectorDataPopulator.setVector( + (IntVector) expectedDelta.getRunEndsVector(), 1, 3, 4, 6); + expectedDelta.setValueCount(6); + assertVectorsEqual(expectedDelta, delta); + } + } + } + @Test public void testAppendVectorNegative() { final int vectorLength = 10;