diff --git a/.asf.yaml b/.asf.yaml index bf17bd0ff6..13bd59c224 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -19,7 +19,7 @@ github: description: "Official Java implementation of Apache Arrow" - homepage: https://arrow.apache.org/ + homepage: https://arrow.apache.org/java/ labels: - apache-arrow - java @@ -29,12 +29,17 @@ github: rebase: false squash: true features: + discussions: true issues: true protected_branches: main: required_linear_history: true notifications: commits: commits@arrow.apache.org + discussions: user@arrow.apache.org issues_status: issues@arrow.apache.org issues_comment: github@arrow.apache.org pullrequests: github@arrow.apache.org +publish: + whoami: asf-site + subdir: java diff --git a/.cmake-format.py b/.cmake-format.py new file mode 100644 index 0000000000..b8fc893969 --- /dev/null +++ b/.cmake-format.py @@ -0,0 +1,76 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# cmake-format configuration file +# Use `archery lint --cmake-format --fix` to reformat all cmake files in the +# source tree + +# ----------------------------- +# Options affecting formatting. +# ----------------------------- +with section("format"): + # How wide to allow formatted cmake files + line_width = 90 + + # How many spaces to tab for indent + tab_size = 2 + + # If a positional argument group contains more than this many arguments, + # then force it to a vertical layout. + max_pargs_hwrap = 4 + + # If the statement spelling length (including space and parenthesis) is + # smaller than this amount, then force reject nested layouts. + # This value only comes into play when considering whether or not to nest + # arguments below their parent. If the number of characters in the parent + # is less than this value, we will not nest. + min_prefix_chars = 32 + + # If true, separate flow control names from their parentheses with a space + separate_ctrl_name_with_space = False + + # If true, separate function names from parentheses with a space + separate_fn_name_with_space = False + + # If a statement is wrapped to more than one line, than dangle the closing + # parenthesis on it's own line + dangle_parens = False + + # What style line endings to use in the output. + line_ending = 'unix' + + # Format command names consistently as 'lower' or 'upper' case + command_case = 'lower' + + # Format keywords consistently as 'lower' or 'upper' case + keyword_case = 'unchanged' + +# ------------------------------------------------ +# Options affecting comment reflow and formatting. +# ------------------------------------------------ +with section("markup"): + # enable comment markup parsing and reflow + enable_markup = False + + # If comment markup is enabled, don't reflow the first comment block in + # eachlistfile. Use this to preserve formatting of your + # copyright/licensestatements. + first_comment_is_literal = True + + # If comment markup is enabled, don't reflow any comment block which + # matches this (regex) pattern. Default is `None` (disabled). + literal_comment_pattern = None diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000000..1b04cdad3c --- /dev/null +++ b/.editorconfig @@ -0,0 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[*.sh] +indent_style = space +indent_size = 2 diff --git a/.env b/.env index bb9f63ab30..ef0c5fb101 100644 --- a/.env +++ b/.env @@ -40,11 +40,17 @@ ARCH_SHORT=amd64 # Default repository to pull and push images from REPO=ghcr.io/apache/arrow-java-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 -MAVEN=3.9.6 +JDK=17 +MAVEN=3.9.9 + +# Versions for various dependencies used to build artifacts +# Keep in sync with apache/arrow +ARROW_REPO_ROOT=./arrow +VCPKG="9b965a116838c6cdcd36bca60d1b81b030c8ab8d" # 2026.05.27 (not release, upstream commit) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000000..8201a5a915 --- /dev/null +++ b/.github/CODEOWNERS @@ -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. +# +# Any committer can add themselves to any of the path patterns +# and will subsequently get requested as a reviewer for any PRs +# that change matching files. +# +# This file uses .gitignore syntax with a few exceptions see the +# documentation about the syntax: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners + +# modules +/adapter/ @lidavidm @laurentgo @wgtmac @jbonofre +/algorithm/ @lidavidm @laurentgo @wgtmac @jbonofre +/arrow-format/ @lidavidm @laurentgo @wgtmac @jbonofre +/c/ @lidavidm @laurentgo @wgtmac @jbonofre +/compression/ @lidavidm @laurentgo @wgtmac @jbonofre +/dataset/ @lidavidm @laurentgo @wgtmac @jbonofre +/flight/ @lidavidm @laurentgo @wgtmac @jbonofre +/format/ @lidavidm @laurentgo @wgtmac @jbonofre +/gandiva/ @lidavidm @laurentgo @wgtmac @jbonofre +/memory/ @lidavidm @laurentgo @wgtmac @jbonofre +/performance/ @lidavidm @laurentgo @wgtmac @jbonofre +/tools/ @lidavidm @laurentgo @wgtmac @jbonofre +/vector/ @lidavidm @laurentgo @wgtmac @jbonofre +CMakeLists.txt @lidavidm @laurentgo @wgtmac @jbonofre + +# release scripts +/ci/ @lidavidm @laurentgo @wgtmac @kou @jbonofre +/dev/ @lidavidm @laurentgo @wgtmac @kou @jbonofre + +# PR CI and repository files +/.github/ @lidavidm @laurentgo @wgtmac @kou @jbonofre +/.asf.yaml @lidavidm @laurentgo @wgtmac @kou @jbonofre +/.env @lidavidm @laurentgo @wgtmac @jbonofre +/.pre-commit-config.yaml @lidavidm @laurentgo @wgtmac @jbonofre +/Brewfile @lidavidm @laurentgo @wgtmac @jbonofre +/docker-compose.yaml @lidavidm @laurentgo @wgtmac @jbonofre + +# Java specific +/.mvn/ @lidavidm @laurentgo @wgtmac @jbonofre +/bom/ @lidavidm @laurentgo @wgtmac @jbonofre +pom.xml @lidavidm @laurentgo @wgtmac @jbonofre diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000000..2ae17814a1 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,7 @@ +## What's Changed + +Please fill in a description of the changes here. + +**This contains breaking changes.** + +Closes #NNN. diff --git a/.github/release.yml b/.github/release.yml new file mode 100644 index 0000000000..f4a1bc9da3 --- /dev/null +++ b/.github/release.yml @@ -0,0 +1,34 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +changelog: + categories: + - title: Breaking Changes + labels: + - breaking-change + + - title: New Features and Enhancements + labels: + - enhancement + + - title: Bug Fixes + labels: + - bug-fix + + - title: Other Changes + labels: + - "*" 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 2ac230f635..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@v4 + 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.js b/.github/workflows/dev_pr.js new file mode 100644 index 0000000000..13acc946e1 --- /dev/null +++ b/.github/workflows/dev_pr.js @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +async function have_comment(github, context, pr_number, tag) { + console.log(`Looking for existing comment on ${pr_number} with substring ${tag}`); + const query = ` +query($owner: String!, $name: String!, $number: Int!, $cursor: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + id + comments (after:$cursor, first: 50) { + nodes { + id + body + author { + login + } + } + pageInfo { + endCursor + hasNextPage + } + } + } + } +}`; + const tag_regexp = new RegExp(tag); + + let cursor = null; + let pr_id = null; + while (true) { + const result = await github.graphql(query, { + owner: context.repo.owner, + name: context.repo.repo, + number: pr_number, + cursor, + }); + pr_id = result.repository.pullRequest.id; + cursor = result.repository.pullRequest.comments.pageInfo; + const comments = result.repository.pullRequest.comments.nodes; + + for (const comment of comments) { + console.log(comment); + if (comment.author.login === "github-actions" && + comment.body.match(tag_regexp) !== null) { + console.log("Found existing comment"); + return {pr_id, comment_id: comment.id}; + } + } + + if (!result.repository.pullRequest.comments.hasNextPage || + comments.length === 0) { + break; + } + } + console.log("No existing comment"); + return {pr_id, comment_id: null}; +} + +async function upsert_comment(github, {pr_id, comment_id}, body, visible) { + console.log(`Upsert comment (pr_id=${pr_id}, comment_id=${comment_id}, visible=${visible})`); + if (!visible) { + if (comment_id === null) return; + + const query = ` +mutation makeComment($comment: ID!) { + minimizeComment(input: {subjectId: $comment, classifier: RESOLVED}) { + clientMutationId + } +}`; + await github.graphql(query, { + comment: comment_id, + body, + }); + return; + } + + if (comment_id === null) { + const query = ` +mutation makeComment($pr: ID!, $body: String!) { + addComment(input: {subjectId: $pr, body: $body}) { + clientMutationId + } +}`; + await github.graphql(query, { + pr: pr_id, + body, + }); + } else { + const query = ` +mutation makeComment($comment: ID!, $body: String!) { + unminimizeComment(input: {subjectId: $comment}) { + clientMutationId + } + updateIssueComment(input: {id: $comment, body: $body}) { + clientMutationId + } +}`; + await github.graphql(query, { + comment: comment_id, + body, + }); + } +} + +module.exports = { + check_title_format: function({core, github, context}) { + const title = context.payload.pull_request.title; + if (title.startsWith("MINOR: ")) { + console.log("PR is a minor PR"); + return {"issue": null}; + } + + const match = title.match(/^GH-([0-9]+): .*$/); + if (match === null) { + core.setFailed("Invalid PR title format. Must either be MINOR: or GH-NNN:"); + return {"issue": null}; + } + return {"issue": parseInt(match[1], 10)}; + }, + + apply_labels: async function({core, github, context}) { + const body = (context.payload.pull_request.body || "").split(/\n/g); + var has_breaking = false; + for (const line of body) { + if (line.trim().startsWith("**This contains breaking changes.**")) { + has_breaking = true; + break; + } + } + if (has_breaking) { + console.log("PR has breaking changes"); + await github.rest.issues.addLabels({ + issue_number: context.payload.pull_request.number, + owner: context.repo.owner, + repo: context.repo.repo, + labels: ["breaking-change"], + }); + } else { + console.log("PR has no breaking changes"); + } + }, + + check_labels: async function({core, github, context}) { + const categories = ["bug-fix", "chore", "dependencies", "documentation", "enhancement"]; + const labels = (context.payload.pull_request.labels || []); + const required = new Set(categories); + var found = false; + + for (const label of labels) { + console.log(`Found label ${label.name}`); + if (required.has(label.name)) { + found = true; + break; + } + } + + // Look to see if we left a comment before. + const comment_tag = "label_helper_comment"; + const maybe_comment_id = await have_comment(github, context, context.payload.pull_request.number, comment_tag); + console.log("Found comment?"); + console.log(maybe_comment_id); + const body_text = ` + +Thank you for opening a pull request! + +Please label the PR with one or more of: + +${categories.map(c => `- ${c}`).join("\n")} + +Also, add the 'breaking-change' label if appropriate. + +See [CONTRIBUTING.md](https://github.com/apache/arrow-java/blob/main/CONTRIBUTING.md) for details. +`; + + if (found) { + console.log("PR has appropriate label(s)"); + await upsert_comment(github, maybe_comment_id, body_text, false); + } else { + console.log(body_text); + await upsert_comment(github, maybe_comment_id, body_text, true); + core.setFailed("Missing required labels. See CONTRIBUTING.md"); + } + }, + + check_linked_issue: async function({core, github, context, issue}) { + console.log(issue); + if (issue.issue === null) { + console.log("This is a MINOR PR"); + return; + } + const expected = `https://github.com/apache/arrow-java/issues/${issue.issue}`; + + const query = ` +query($owner: String!, $name: String!, $number: Int!) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + closingIssuesReferences(first: 50) { + edges { + node { + number + } + } + } + } + } +}`; + + const result = await github.graphql(query, { + owner: context.repo.owner, + name: context.repo.repo, + number: context.payload.pull_request.number, + }); + const issues = result.repository.pullRequest.closingIssuesReferences.edges; + console.log(issues); + + for (const link of issues) { + console.log(`PR is linked to ${link.node.number}`); + if (link.node.number === issue.issue) { + console.log(`Found link to ${expected}`); + return; + } + } + console.log(`Did not find link to ${expected}`); + core.setFailed("Missing link to issue in title"); + }, +}; diff --git a/.github/workflows/dev_pr.yml b/.github/workflows/dev_pr.yml new file mode 100644 index 0000000000..ad000df88e --- /dev/null +++ b/.github/workflows/dev_pr.yml @@ -0,0 +1,85 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +name: Dev PR + +on: + pull_request_target: + types: + - labeled + - unlabeled + - opened + - edited + - reopened + - synchronize + - ready_for_review + - review_requested + +concurrency: + group: ${{ github.repository }}-${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +permissions: + contents: read + issues: write + pull-requests: write + +jobs: + pr-label: + name: "Ensure PR format" + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Ensure PR title format + id: title-format + 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@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@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@v9 + with: + script: | + const scripts = require(`${process.env.GITHUB_WORKSPACE}/.github/workflows/dev_pr.js`); + await scripts.check_linked_issue({core, github, context, issue: ${{ steps.title-format.outputs.result }}}); + + - name: Assign milestone + if: '! github.event.pull_request.draft' + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + ./.github/workflows/dev_pr_milestone.sh "${GITHUB_REPOSITORY}" ${{ github.event.number }} diff --git a/.github/workflows/dev_pr_milestone.sh b/.github/workflows/dev_pr_milestone.sh new file mode 100755 index 0000000000..4a77eb1f73 --- /dev/null +++ b/.github/workflows/dev_pr_milestone.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Assign a milestone to the given PR based on the open milestones and known +# releases. + +set -euo pipefail + +main() { + local -r repo="${1}" + local -r pr_number="${2}" + echo "On ${repo} pull ${pr_number}" + + local -r existing_milestone=$(gh pr view "${pr_number}" \ + --json milestone \ + -t '{{if .milestone}}{{.milestone.title}}{{end}}') + + if [[ -n "${existing_milestone}" ]]; then + echo "PR has milestone: ${existing_milestone}" + local -r milestone="${existing_milestone}" + else + local -r milestone=$( + gh api "/repos/${repo}/milestones" | + jq --raw-output '.[] | .title' | + grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' | + head -n1 + ) + + echo "Assigning milestone: ${milestone}" + gh pr edit "${pr_number}" -m "${milestone}" + fi + + local -r repo_owner=$(echo "${repo}" | cut -d'/' -f1) + local -r repo_name=$(echo "${repo}" | cut -d'/' -f2) + local -r graphql_query="{ + repository(owner: \"${repo_owner}\", name: \"${repo_name}\") { + pullRequest(number: ${pr_number}) { + closingIssuesReferences(first: 5) { + edges { + node { + number + } + } + } + } + } + }" + local -r linked_issues=$(gh api graphql -f query="${graphql_query}" | jq -r '.data.repository.pullRequest.closingIssuesReferences.edges | .[].node.number') + for issue in ${linked_issues}; do + echo "Linked issue: ${issue}" + gh issue edit "${issue}" --milestone "${milestone}" + done +} + +main "$@" 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 new file mode 100644 index 0000000000..18a721ac02 --- /dev/null +++ b/.github/workflows/rc.yml @@ -0,0 +1,617 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT 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: RC +on: + push: + branches: + - "**" + - "!dependabot/**" + tags: + - "*-rc*" + pull_request: + schedule: + - cron: "0 0 * * *" +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true +permissions: + contents: read +jobs: + source: + name: Source + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Checkout + uses: actions/checkout@v7 + with: + submodules: recursive + - name: Prepare for tag + if: github.ref_type == 'tag' + run: | + version=${GITHUB_REF_NAME%-rc*} + version=${version#v} + rc=${GITHUB_REF_NAME#*-rc} + echo "VERSION=${version}" >> ${GITHUB_ENV} + echo "RC=${rc}" >> ${GITHUB_ENV} + - name: Prepare for branch + if: github.ref_type == 'branch' + run: | + version=$(grep -o '^ .*' "pom.xml" | + sed \ + -e 's,^ ,,' \ + -e 's,$,,') + rc=$(date +%Y%m%d) + echo "VERSION=${version}" >> ${GITHUB_ENV} + echo "RC=${rc}" >> ${GITHUB_ENV} + - name: Archive + run: | + id="apache-arrow-java-${VERSION}" + tar_gz="${id}.tar.gz" + echo "TAR_GZ=${tar_gz}" >> ${GITHUB_ENV} + git archive HEAD --prefix "${id}/" --output "${tar_gz}" + sha256sum "${tar_gz}" > "${tar_gz}.sha256" + sha512sum "${tar_gz}" > "${tar_gz}.sha512" + - name: Audit + run: | + dev/release/run_rat.sh "${TAR_GZ}" + - name: Upload source archive + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: release-source + path: | + apache-arrow-java-* + jni-linux: + name: JNI ${{ matrix.platform.runs_on }} ${{ matrix.platform.arch }} + runs-on: ${{ matrix.platform.runs_on }} + timeout-minutes: 120 + needs: + - source + strategy: + fail-fast: false + matrix: + platform: + - runs_on: ubuntu-latest + arch: "x86_64" + archery_arch: "amd64" + - runs_on: ubuntu-24.04-arm + arch: "aarch_64" + archery_arch: "arm64v8" + env: + # architecture name used for archery build + ARCH: ${{ matrix.platform.archery_arch }} + DOCKER_VOLUME_PREFIX: .docker/ + permissions: + contents: read + packages: write + steps: + - name: Download source archive + 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: Checkout Apache Arrow C++ + if: github.event_name == 'schedule' + uses: actions/checkout@v7 + with: + repository: apache/arrow + path: arrow + - name: Checkout apache/arrow-testing + uses: actions/checkout@v7 + with: + repository: apache/arrow-testing + path: arrow/testing + - name: Checkout apache/parquet-testing + uses: actions/checkout@v7 + with: + repository: apache/parquet-testing + path: arrow/cpp/submodules/parquet-testing + - uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Cache + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: .docker + key: jni-linux-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} + restore-keys: jni-linux-${{ matrix.platform.arch }}- + - name: Build + run: | + docker compose run vcpkg-jni + - name: Push Docker image + if: success() && github.event_name == 'push' && github.repository == 'apache/arrow-java' && github.ref_name == 'main' + run: | + docker compose push vcpkg-jni + - 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@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: 60 + needs: + - source + strategy: + fail-fast: false + matrix: + platform: + - { 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@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: Checkout Apache Arrow C++ + if: github.event_name == 'schedule' + uses: actions/checkout@v7 + with: + repository: apache/arrow + path: arrow + - name: Checkout apache/arrow-testing + uses: actions/checkout@v7 + with: + repository: apache/arrow-testing + path: arrow/testing + - name: Checkout apache/parquet-testing + uses: actions/checkout@v7 + with: + repository: apache/parquet-testing + path: arrow/cpp/submodules/parquet-testing + - name: Set up Python + uses: actions/setup-python@v7 + with: + cache: 'pip' + python-version: 3.12 + - name: Install Archery + run: pip install -e arrow/dev/archery[all] + - name: Install dependencies + run: | + # We want to use llvm@14 to avoid shared z3 + # dependency. llvm@14 doesn't depend on z3 and llvm depends + # on z3. And Homebrew's z3 provides only shared library. It + # doesn't provides static z3 because z3's CMake doesn't accept + # building both shared and static libraries at once. + # See also: Z3_BUILD_LIBZ3_SHARED in + # https://github.com/Z3Prover/z3/blob/master/README-CMake.md + # + # If llvm is installed, Apache Arrow C++ uses llvm rather than + # 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 + # Homebrew exists in /usr/local on GitHub Actions. If + # Homebrew's python@XXX is updated without "--overwrite", it + # 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@ | sort -r); do + brew install --overwrite ${python_package} + done + brew install --overwrite python3 + + if [ "$(uname -m)" = "arm64" ]; then + # pkg-config formula is deprecated but it's still installed + # in GitHub Actions runner now. We can remove this once + # pkg-config formula is removed from GitHub Actions runner. + brew uninstall pkg-config || : + brew uninstall pkg-config@0.29.2 || : + fi + + brew bundle --file=arrow/cpp/Brewfile + # We want to link aws-sdk-cpp statically but Homebrew's + # aws-sdk-cpp provides only shared library. If we have + # Homebrew's aws-sdk-cpp, our build mix Homebrew's + # aws-sdk-cpp and bundled aws-sdk-cpp. We uninstall Homebrew's + # aws-sdk-cpp to ensure using only bundled aws-sdk-cpp. + brew uninstall aws-sdk-cpp + # We want to use bundled RE2 for static linking. If + # Homebrew's RE2 is installed, its header file may be used. + # We uninstall Homebrew's RE2 to ensure using bundled RE2. + brew uninstall grpc || : # gRPC depends on RE2 + brew uninstall grpc@1.54 || : # gRPC 1.54 may be installed too + brew uninstall re2 + # We want to use bundled Protobuf for static linking. If + # Homebrew's Protobuf is installed, its library file may be + # used on test We uninstall Homebrew's Protobuf to ensure using + # bundled Protobuf. + brew uninstall protobuf + + brew bundle --file=Brewfile + - name: Prepare ccache + run: | + echo "CCACHE_DIR=${PWD}/ccache" >> ${GITHUB_ENV} + - name: Cache ccache + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: ccache + key: jni-macos-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} + restore-keys: jni-macos-${{ matrix.platform.arch }}- + - name: Build + run: | + set -e + # make brew Java available to CMake + 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@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: jni-macos-${{ matrix.platform.arch }} + path: jni-macos-${{ matrix.platform.arch }}.tar.gz + jni-windows: + name: JNI ${{ matrix.platform.runs_on }} ${{ matrix.platform.arch }} + runs-on: ${{ matrix.platform.runs_on }} + timeout-minutes: 45 + needs: + - source + strategy: + fail-fast: false + matrix: + platform: + - runs_on: windows-2022 + arch: "x86_64" + steps: + - name: Download source archive + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: release-source + - name: Extract source archive + shell: bash + run: | + tar -xf apache-arrow-java-*.tar.gz --strip-components=1 + - name: Download the latest Apache Arrow C++ + if: github.event_name != 'schedule' + shell: bash + run: | + ci/scripts/download_cpp.sh + - name: Checkout Apache Arrow C++ + if: github.event_name == 'schedule' + uses: actions/checkout@v7 + with: + repository: apache/arrow + path: arrow + - name: Set up Java + uses: actions/setup-java@v5 + with: + java-version: '17' + distribution: 'temurin' + - name: Download Timezone Database + shell: bash + run: | + arrow/ci/scripts/download_tz_database.sh + - name: Install ccache + shell: bash + run: | + env | sort + version=4.10.2 + base_name="ccache-${version}-windows-x86_64" + url="https://github.com/ccache/ccache/releases/download/v${version}/${base_name}.zip" + curl --fail --location --remote-name "${url}" + unzip "${base_name}.zip" + chmod +x "${base_name}/ccache.exe" + mv "${base_name}/ccache.exe" /usr/bin/ + rm -rf "${base_name}"{,.zip} + - name: Prepare ccache + shell: bash + run: | + echo "CCACHE_DIR=${PWD}/ccache" >> ${GITHUB_ENV} + - name: Cache ccache + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: ccache + key: jni-windows-${{ matrix.platform.arch }}-${{ hashFiles('arrow/cpp/**') }} + restore-keys: jni-windows-${{ matrix.platform.arch }}- + - name: Build + shell: cmd + run: | + 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" + - name: Compress into single artifact to keep directory structure + shell: bash + run: tar -cvzf jni-windows-${{ matrix.platform.arch }}.tar.gz jni/ + - name: Upload artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: jni-windows-${{ matrix.platform.arch }} + path: jni-windows-${{ matrix.platform.arch }}.tar.gz + binaries: + name: Binaries + runs-on: ubuntu-latest + needs: + - jni-linux + - jni-macos + - jni-windows + steps: + - name: Download artifacts + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + path: artifacts + - name: Decompress artifacts + run: | + mv artifacts/*/*.tar.gz . + tar -xf apache-arrow-java-*.tar.gz --strip-components=1 + tar -xvzf jni-linux-x86_64.tar.gz + tar -xvzf jni-linux-aarch_64.tar.gz + tar -xvzf jni-macos-x86_64.tar.gz + tar -xvzf jni-macos-aarch_64.tar.gz + tar -xvzf jni-windows-x86_64.tar.gz + - name: Test that shared libraries exist + run: | + set -x + + test -f jni/arrow_cdata_jni/x86_64/libarrow_cdata_jni.so + test -f jni/arrow_dataset_jni/x86_64/libarrow_dataset_jni.so + test -f jni/arrow_orc_jni/x86_64/libarrow_orc_jni.so + test -f jni/gandiva_jni/x86_64/libgandiva_jni.so + + test -f jni/arrow_cdata_jni/aarch_64/libarrow_cdata_jni.so + test -f jni/arrow_dataset_jni/aarch_64/libarrow_dataset_jni.so + test -f jni/arrow_orc_jni/aarch_64/libarrow_orc_jni.so + test -f jni/gandiva_jni/aarch_64/libgandiva_jni.so + + test -f jni/arrow_cdata_jni/x86_64/libarrow_cdata_jni.dylib + test -f jni/arrow_dataset_jni/x86_64/libarrow_dataset_jni.dylib + test -f jni/arrow_orc_jni/x86_64/libarrow_orc_jni.dylib + test -f jni/gandiva_jni/x86_64/libgandiva_jni.dylib + + test -f jni/arrow_cdata_jni/aarch_64/libarrow_cdata_jni.dylib + test -f jni/arrow_dataset_jni/aarch_64/libarrow_dataset_jni.dylib + test -f jni/arrow_orc_jni/aarch_64/libarrow_orc_jni.dylib + test -f jni/gandiva_jni/aarch_64/libgandiva_jni.dylib + + test -f jni/arrow_cdata_jni/x86_64/arrow_cdata_jni.dll + 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@v7 + with: + repository: apache/arrow-testing + path: testing + - name: Cache ~/.m2 + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: ~/.m2 + key: binaries-build-${{ hashFiles('**/*.java', '**/pom.xml') }} + restore-keys: binaries-build- + - name: Build bundled JAR and docs + run: | + ci/scripts/jni_full_build.sh . jni binaries + - name: Prepare docs + run: | + mkdir -p docs + cp -a target/site/apidocs reference + tar -cvzf reference.tar.gz reference + - name: Upload binaries + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: release-binaries + path: binaries/* + - name: Upload docs + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: reference + path: reference.tar.gz + docs: + name: Docs + needs: + - binaries + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + steps: + - uses: actions/setup-python@v7 + with: + cache: 'pip' + - name: Download source archive + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: release-source + - name: Download Javadocs + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: reference + - name: Extract source archive + run: | + tar -xf apache-arrow-java-*.tar.gz --strip-components=1 + - name: Build + run: | + cd docs + python -m venv venv + source venv/bin/activate + pip install -r requirements.txt + make html + tar -xf ../reference.tar.gz -C build/html + - 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@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: release-html + path: html.tar.gz + publish-docs: + name: Publish docs + # Run only when: + # * We push to a branch + # * If the target repository is apache/arrow-java: + # * The target branch is main + # * Else (fork repositories): + # * All branches + # * We can preview the last pushed content + # at https://${YOUR_GITHUB_ACCOUNT}.github.io/arrow-java/main/ + if: >- + github.event_name == 'push' && + github.ref_type == 'branch' && + ((github.repository == 'apache/arrow-java' && github.ref_name == 'main') || + github.repository != 'apache/arrow-java') + needs: + - docs + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Checkout + uses: actions/checkout@v7 + with: + path: site + - name: Prepare branch + run: | + if [ "${GITHUB_REPOSITORY}" = "apache/arrow-java" ]; then + BRANCH=asf-site + else + BRANCH=gh-pages + fi + echo "BRANCH=${BRANCH}" >> ${GITHUB_ENV} + + cp site/.asf.yaml ./ + cd site + git fetch + if ! git switch -c "${BRANCH}" "origin/${BRANCH}"; then + git switch --orphan "${BRANCH}" + fi + touch .nojekyll + cp ../.asf.yaml ./ + git add .nojekyll .asf.yaml + - name: Download + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: release-html + - name: Extract + run: | + tar -xf html.tar.gz + rm -rf site/main + mv html site/main + git -C site add main + if [ "$(git status --prcelain)" == "" ]; then + NEED_PUSH=true + else + NEED_PUSH=false + fi + echo "NEED_PUSH=${NEED_PUSH}" >> ${GITHUB_ENV} + - name: Push + if: env.NEED_PUSH == 'true' + run: | + cd site + git config --global user.name 'github-actions[bot]' + git config --global user.email 'github-actions[bot]@users.noreply.github.com' + git commit -m "Publish documentation (${GITHUB_SHA})" + git push origin "${BRANCH}" + verify: + name: Verify + needs: + - binaries + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: + - macos-latest + - ubuntu-latest + steps: + - name: Download release artifacts + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: release-* + - name: Verify + run: | + mv release-source/* ./ + tar_gz=$(echo apache-arrow-java-*.tar.gz) + version=${tar_gz#apache-arrow-java-} + version=${version%.tar.gz} + # rc isn't used with VERIFY_DOWNLOAD=0 + if [ "${GITHUB_REF_TYPE}" = "tag" ]; then + rc="${GITHUB_REF_NAME#*-rc}" + else + rc=$(date +%Y%m%d) + fi + tar -xf ${tar_gz} + export VERIFY_DEFAULT=0 + export VERIFY_BINARY=1 + export VERIFY_SOURCE=1 + cd apache-arrow-java-${version} + mv ../${tar_gz}* ./ + mv ../release-binaries binaries + dev/release/verify_rc.sh "${version}" "${rc}" + upload: + name: Upload + if: github.ref_type == 'tag' + needs: + - docs + - verify + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - name: Download release artifacts + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: release-* + path: artifacts + - name: Upload + run: | + # GH-499: How to create release notes? + version=${GITHUB_REF_NAME%-rc*} + version=${version#v} + rc=${GITHUB_REF_NAME#*-rc} + gh release create ${GITHUB_REF_NAME} \ + --generate-notes \ + --prerelease \ + --repo ${GITHUB_REPOSITORY} \ + --title "Apache Arrow Java ${version} RC${rc}" \ + --verify-tag + # GitHub CLI does not respect their own rate limits + # https://github.com/cli/cli/issues/9586 + for artifact in artifacts/*/*; do + sleep 1 + gh release upload ${GITHUB_REF_NAME} \ + --repo ${GITHUB_REPOSITORY} \ + $artifact + done + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000000..7692eb6cbe --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,96 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT 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: Release +on: + push: + tags: + - "*" + - "!*-rc*" +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true +permissions: + contents: write +env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} +jobs: + publish: + name: Publish + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Download RC contents + run: | + set -x + latest_rc_tag=$(gh release list \ + --jq '.[].tagName' \ + --json tagName \ + --repo ${GITHUB_REPOSITORY} | \ + grep -F "${GITHUB_REF_NAME}-rc" | \ + head -n1) + gh release download ${latest_rc_tag} \ + --repo ${GITHUB_REPOSITORY} \ + --dir dists + - name: Create GitHub Release + run: | + # GH-499: How to create release notes? + version=${GITHUB_REF_NAME#v} + gh release create ${GITHUB_REF_NAME} \ + --generate-notes \ + --repo ${GITHUB_REPOSITORY} \ + --title "Apache Arrow Java ${version}" \ + --verify-tag + + # GitHub CLI does not respect their own rate limits + # https://github.com/cli/cli/issues/9586 + for artifact in dists/*; do + sleep 1 + gh release upload ${GITHUB_REF_NAME} \ + --repo ${GITHUB_REPOSITORY} \ + $artifact + done + - name: Checkout for publishing docs + uses: actions/checkout@v7 + with: + path: site + - name: Publish docs + run: | + set -x + + tar -xf dists/html.tar.gz + version=${GITHUB_REF_NAME#v} + + if [ "${GITHUB_REPOSITORY}" = "apache/arrow-java" ]; then + BRANCH=asf-site + else + BRANCH=gh-pages + fi + + cd site + git fetch + git switch -c "${BRANCH}" "origin/${BRANCH}" + + rm -rf current ${version} + cp -a ../html current + cp -a ../html ${version} + git add current ${version} + + git config --global user.name 'github-actions[bot]' + git config --global user.email 'github-actions[bot]@users.noreply.github.com' + git commit -m "Publish documentation (${GITHUB_REF_NAME})" + git push origin "${BRANCH}" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 473ce84ca7..653b16fa32 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -38,35 +38,40 @@ env: jobs: ubuntu: - name: AMD64 Ubuntu 22.04 JDK ${{ matrix.jdk }} Maven ${{ matrix.maven }} + name: AMD64 ${{ matrix.name }} JDK ${{ matrix.jdk }} Maven ${{ matrix.maven }} runs-on: ubuntu-latest if: ${{ !contains(github.event.pull_request.title, 'WIP') }} timeout-minutes: 30 strategy: fail-fast: false matrix: - jdk: [11, 17, 21, 22] - maven: [3.9.6] - image: [java] + jdk: [17, 21, 23] + maven: [3.9.9] + image: [ubuntu, conda-jni-cdata] + include: + - image: ubuntu + name: "Ubuntu" + - image: conda-jni-cdata + name: "Conda JNI" env: JDK: ${{ matrix.jdk }} 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@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2 + uses: actions/cache@v6 with: path: .docker - key: maven-${{ matrix.jdk }}-${{ matrix.maven }}-${{ hashFiles('**/docker-compose.yml', '**/pom.xml') }} + key: maven-${{ matrix.jdk }}-${{ matrix.maven }}-${{ hashFiles('compose.yaml', '**/pom.xml') }} restore-keys: maven-${{ matrix.jdk }}-${{ matrix.maven }}- - name: Execute Docker Build env: # Enables build caching, but not strictly required - DEVELOCITY_ACCESS_KEY: ${{ secrets.GE_ACCESS_TOKEN }} + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} run: | docker compose run \ -e CI=true \ @@ -83,32 +88,33 @@ 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: - DEVELOCITY_ACCESS_KEY: ${{ secrets.GE_ACCESS_TOKEN }} - run: ci/scripts/java_build.sh $(pwd) $(pwd)/build + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + run: ci/scripts/build.sh . build jni - name: Test shell: bash env: - DEVELOCITY_ACCESS_KEY: ${{ secrets.GE_ACCESS_TOKEN }} - run: ci/scripts/java_test.sh $(pwd) $(pwd)/build + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + run: ci/scripts/test.sh . build jni windows: name: AMD64 Windows Server 2022 Java JDK ${{ matrix.jdk }} @@ -118,25 +124,26 @@ 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: - DEVELOCITY_ACCESS_KEY: ${{ secrets.GE_ACCESS_TOKEN }} - run: ci/scripts/java_build.sh $(pwd) $(pwd)/build + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + run: ci/scripts/build.sh . build jni - name: Test shell: bash env: - DEVELOCITY_ACCESS_KEY: ${{ secrets.GE_ACCESS_TOKEN }} - run: ci/scripts/java_test.sh $(pwd) $(pwd)/build + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + run: ci/scripts/test.sh . build jni diff --git a/.gitignore b/.gitignore index ca0ac32463..17d1d43ae1 100644 --- a/.gitignore +++ b/.gitignore @@ -7,16 +7,22 @@ .buildpath .checkstyle .classpath +.cursor/ .factorypath .idea/ .project .settings/ +.vscode/ /*-build/ /.mvn/.develocity/ +/apache-arrow-java-* +/apache-arrow-java.tar.gz /build/ +/dev/release/.env /dev/release/apache-rat-0.16.1.jar /dev/release/filtered_rat.txt /dev/release/rat.xml +/docs/build/ CMakeCache.txt CMakeFiles/ Makefile diff --git a/.mvn/develocity.xml b/.mvn/develocity.xml index df3cbccd2b..298f1efcff 100644 --- a/.mvn/develocity.xml +++ b/.mvn/develocity.xml @@ -20,19 +20,18 @@ --> + arrow - https://ge.apache.org + https://develocity.apache.org false - - true - true - true - #{isFalse(env['CI'])} - true - true + + + + + #{{'0.0.0.0'}} diff --git a/.mvn/extensions.xml b/.mvn/extensions.xml index 0836fc47d0..74482cb2c4 100644 --- a/.mvn/extensions.xml +++ b/.mvn/extensions.xml @@ -23,11 +23,11 @@ com.gradle develocity-maven-extension - 1.22.2 + 2.5.0 com.gradle common-custom-user-data-maven-extension - 2.0.1 + 2.3.0 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 63b7ab9f91..8e8faf3cc0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,16 +16,20 @@ # under the License. repos: + - repo: https://github.com/cheshirekow/cmake-format-precommit + rev: v0.6.13 + hooks: + - id: cmake-format - repo: https://github.com/pre-commit/pre-commit-hooks rev: cef0300fd0fc4d2a87a85fa2093c6b283ea36f4b # v5.0.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer + - id: check-shebang-scripts-are-executable - id: check-yaml - id: check-added-large-files - id: file-contents-sorter files: .gitignore - - repo: local hooks: - id: rat @@ -34,8 +38,23 @@ repos: entry: | bash -c " \ git archive HEAD \ - --prefix=apache-arrow-go/ \ - --output=apache-arrow-go.tar.gz && \ - dev/release/run_rat.sh apache-arrow-go.tar.gz" + --prefix=apache-arrow-java/ \ + --output=apache-arrow-java.tar.gz && \ + dev/release/run_rat.sh apache-arrow-java.tar.gz && \ + rm -f apache-arrow-java.tar.gz" always_run: true pass_filenames: false + - repo: https://github.com/koalaman/shellcheck-precommit + rev: v0.10.0 + hooks: + - id: shellcheck + args: + - "--external-sources" + - repo: https://github.com/scop/pre-commit-shfmt + rev: v3.9.0-1 + hooks: + - id: shfmt + args: + # The default args is "--write --simplify" but we don't use + # "--simplify". Because it's conflicted will ShellCheck. + - "--write" 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/CMakeLists.txt b/CMakeLists.txt index 8b29f37d80..318bd4d10c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -66,7 +66,10 @@ add_library(jni INTERFACE IMPORTED) set_target_properties(jni PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${JNI_INCLUDE_DIRS}") include(CTest) -if(BUILD_TESTING) +if(BUILD_TESTING + AND (ARROW_JAVA_JNI_ENABLE_DATASET + OR ARROW_JAVA_JNI_ENABLE_GANDIVA + OR ARROW_JAVA_JNI_ENABLE_ORC)) find_package(ArrowTesting REQUIRED) find_package(GTest REQUIRED) add_library(arrow_java_test INTERFACE IMPORTED) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8388b1d6c7..680750070f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -30,8 +30,36 @@ existing Arrow issues in [GitHub](https://github.com/apache/arrow-java/issues). ## Did you write a patch that fixes a bug or brings an improvement? -Create a GitHub issue and submit your changes as a GitHub Pull Request. -Please make sure to [reference the issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword) in your PR description. +- Create a GitHub issue and submit your changes as a GitHub Pull Request. +- [Reference the issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword) in your PR description. +- Add one or more of the labels "bug-fix", "chore", "dependencies", "documentation", and "enhancement" to your PR as appropriate. + - "bug-fix" is for PRs that fix a bug. + - "chore" is for other administrative work (build system, release process, etc.). + - "dependencies" is for PRs that upgrade a dependency. (Usually only used by dependabot.) + - "documentation" is for documentation updates. + - "enhancement" is for PRs that add new features. +- Add the "breaking-change" label to your PR if there are breaking API changes. +- Add the PR title. The PR title will be used as the eventual commit message, so please make it descriptive but succinct. + +Example #1: + +``` +GH-12345: Document the pull request process + +Explain how to open a pull request and what the title, body, and labels should be. + +Closes #12345. +``` + +Example #2: + +``` +GH-42424: Expose Netty server builder in Flight + +Allow direct usage of gRPC APIs for low-level control. + +Closes #42424. +``` ### Minor Fixes diff --git a/LICENSE.txt b/LICENSE.txt index 7ae0e2080e..f2e413b63d 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -202,16 +202,21 @@ limitations under the License. -------------------------------------------------------------------------------- -vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java -vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java +This product includes code from Netty 4.1.117.Final: -These file are derived from code from Netty, which is made available under the -Apache License 2.0. +* vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java +* vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java +* memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java + +Copyright: 2014 The Netty Project +Home page: https://netty.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 -------------------------------------------------------------------------------- -memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java +This product includes code from Google Guava 33.4.0-jre + +* memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java -This product includes software from Google Guava, which is made available under the -Apache License 2.0. - * Copyright (C) 2007 The Guava Authors - * https://github.com/google/guava +Copyright: (C) 2011 The Guava Authors +Home page: https://github.com/google/guava/ +License: https://www.apache.org/licenses/LICENSE-2.0 diff --git a/NOTICE.txt b/NOTICE.txt index a595306b89..fd8dd2802a 100644 --- a/NOTICE.txt +++ b/NOTICE.txt @@ -1,5 +1,272 @@ Apache Arrow Java -Copyright 2016-2024 The Apache Software Foundation +Copyright 2016-2025 The Apache Software Foundation This product includes software developed at The Apache Software Foundation (http://www.apache.org/). + +--------------------------------------------------- +This product includes code from Netty 4.1.117.Final, with the following in its NOTICE: + +| The Netty Project +| ================= +| +| Please visit the Netty web site for more information: +| +| * https://netty.io/ +| +| Copyright 2014 The Netty Project +| +| The Netty Project licenses this file to you under the Apache License, +| version 2.0 (the "License"); you may not use this file except in compliance +| with the License. You may obtain a copy of the License at: +| +| https://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +| WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +| License for the specific language governing permissions and limitations +| under the License. +| +| Also, please refer to each LICENSE..txt file, which is located in +| the 'license' directory of the distribution file, for the license terms of the +| components that this product depends on. +| +| ------------------------------------------------------------------------------- +| This product contains the extensions to Java Collections Framework which has +| been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene: +| +| * LICENSE: +| * license/LICENSE.jsr166y.txt (Public Domain) +| * HOMEPAGE: +| * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/ +| * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/ +| +| This product contains a modified version of Robert Harder's Public Domain +| Base64 Encoder and Decoder, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.base64.txt (Public Domain) +| * HOMEPAGE: +| * http://iharder.sourceforge.net/current/java/base64/ +| +| This product contains a modified portion of 'Webbit', an event based +| WebSocket and HTTP server, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.webbit.txt (BSD License) +| * HOMEPAGE: +| * https://github.com/joewalnes/webbit +| +| This product contains a modified portion of 'SLF4J', a simple logging +| facade for Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.slf4j.txt (MIT License) +| * HOMEPAGE: +| * https://www.slf4j.org/ +| +| This product contains a modified portion of 'Apache Harmony', an open source +| Java SE, which can be obtained at: +| +| * NOTICE: +| * license/NOTICE.harmony.txt +| * LICENSE: +| * license/LICENSE.harmony.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://archive.apache.org/dist/harmony/ +| +| This product contains a modified portion of 'jbzip2', a Java bzip2 compression +| and decompression library written by Matthew J. Francis. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jbzip2.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jbzip2/ +| +| This product contains a modified portion of 'libdivsufsort', a C API library to construct +| the suffix array and the Burrows-Wheeler transformed string for any input string of +| a constant-size alphabet written by Yuta Mori. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.libdivsufsort.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/y-256/libdivsufsort +| +| This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM, +| which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jctools.txt (ASL2 License) +| * HOMEPAGE: +| * https://github.com/JCTools/JCTools +| +| This product optionally depends on 'JZlib', a re-implementation of zlib in +| pure Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jzlib.txt (BSD style License) +| * HOMEPAGE: +| * http://www.jcraft.com/jzlib/ +| +| This product optionally depends on 'Compress-LZF', a Java library for encoding and +| decoding data in LZF format, written by Tatu Saloranta. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.compress-lzf.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/ning/compress +| +| This product optionally depends on 'lz4', a LZ4 Java compression +| and decompression library written by Adrien Grand. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lz4.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jpountz/lz4-java +| +| This product optionally depends on 'lzma-java', a LZMA Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lzma-java.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jponge/lzma-java +| +| This product optionally depends on 'zstd-jni', a zstd-jni Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.zstd-jni.txt (BSD) +| * HOMEPAGE: +| * https://github.com/luben/zstd-jni +| +| This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression +| and decompression library written by William Kinney. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jfastlz.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jfastlz/ +| +| This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data +| interchange format, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.protobuf.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/protobuf +| +| This product optionally depends on 'Bouncy Castle Crypto APIs' to generate +| a temporary self-signed X.509 certificate when the JVM does not provide the +| equivalent functionality. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.bouncycastle.txt (MIT License) +| * HOMEPAGE: +| * https://www.bouncycastle.org/ +| +| This product optionally depends on 'Snappy', a compression library produced +| by Google Inc, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.snappy.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/snappy +| +| This product optionally depends on 'JBoss Marshalling', an alternative Java +| serialization API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jboss-marshalling.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jboss-remoting/jboss-marshalling +| +| This product optionally depends on 'Caliper', Google's micro- +| benchmarking framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.caliper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/google/caliper +| +| This product optionally depends on 'Apache Commons Logging', a logging +| framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-logging.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/logging/ +| +| This product optionally depends on 'Apache Log4J', a logging framework, which +| can be obtained at: +| +| * LICENSE: +| * license/LICENSE.log4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://logging.apache.org/log4j/ +| +| This product optionally depends on 'Aalto XML', an ultra-high performance +| non-blocking XML processor, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.aalto-xml.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://wiki.fasterxml.com/AaltoHome +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hpack.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/twitter/hpack +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Cory Benfield. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hyper-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/python-hyper/hpack/ +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Tatsuhiro Tsujikawa. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.nghttp2-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/nghttp2/nghttp2/ +| +| This product contains a modified portion of 'Apache Commons Lang', a Java library +| provides utilities for the java.lang API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-lang.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/proper/commons-lang/ +| +| +| This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build. +| +| * LICENSE: +| * license/LICENSE.mvn-wrapper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/takari/maven-wrapper +| +| This product contains the dnsinfo.h header file, that provides a way to retrieve the system DNS configuration on MacOS. +| This private header is also used by Apple's open source +| mDNSResponder (https://opensource.apple.com/tarballs/mDNSResponder/). +| +| * LICENSE: +| * license/LICENSE.dnsinfo.txt (Apple Public Source License 2.0) +| * HOMEPAGE: +| * https://www.opensource.apple.com/source/configd/configd-453.19/dnsinfo/dnsinfo.h +| +| This product optionally depends on 'Brotli4j', Brotli compression and +| decompression for Java., which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.brotli4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/hyperxpro/Brotli4j diff --git a/README.md b/README.md index 9f1b1c63c8..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 @@ -48,10 +48,10 @@ a version of your choosing. ```bash $ flatc --version -flatc version 24.3.25 +flatc version 25.1.24 -$ grep "dep.fbs.version" java/pom.xml - 24.3.25 +$ grep "dep.fbs.version" pom.xml + 25.1.24 ``` 2. Generate the flatbuffer java files by performing the following: @@ -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/module-info.java b/adapter/avro/src/main/java/module-info.java index 5c6204be60..fee6c72199 100644 --- a/adapter/avro/src/main/java/module-info.java +++ b/adapter/avro/src/main/java/module-info.java @@ -18,6 +18,8 @@ module org.apache.arrow.adapter.avro { exports org.apache.arrow.adapter.avro.consumers; exports org.apache.arrow.adapter.avro.consumers.logical; + exports org.apache.arrow.adapter.avro.producers; + exports org.apache.arrow.adapter.avro.producers.logical; exports org.apache.arrow.adapter.avro; requires org.apache.arrow.memory.core; 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 new file mode 100644 index 0000000000..e09b99f670 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java @@ -0,0 +1,694 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +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; +import org.apache.arrow.adapter.avro.producers.AvroFloat4Producer; +import org.apache.arrow.adapter.avro.producers.AvroFloat8Producer; +import org.apache.arrow.adapter.avro.producers.AvroIntProducer; +import org.apache.arrow.adapter.avro.producers.AvroListProducer; +import org.apache.arrow.adapter.avro.producers.AvroMapProducer; +import org.apache.arrow.adapter.avro.producers.AvroNullProducer; +import org.apache.arrow.adapter.avro.producers.AvroNullableProducer; +import org.apache.arrow.adapter.avro.producers.AvroSmallIntProducer; +import org.apache.arrow.adapter.avro.producers.AvroStringProducer; +import org.apache.arrow.adapter.avro.producers.AvroStructProducer; +import org.apache.arrow.adapter.avro.producers.AvroTinyIntProducer; +import org.apache.arrow.adapter.avro.producers.AvroUint1Producer; +import org.apache.arrow.adapter.avro.producers.AvroUint2Producer; +import org.apache.arrow.adapter.avro.producers.AvroUint4Producer; +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; +import org.apache.arrow.adapter.avro.producers.logical.AvroDecimal256Producer; +import org.apache.arrow.adapter.avro.producers.logical.AvroDecimalProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimeMicroProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimeMilliProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimeNanoProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimeSecProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampMicroProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampMicroTzProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampMilliProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampMilliTzProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampNanoProducer; +import org.apache.arrow.adapter.avro.producers.logical.AvroTimestampNanoTzProducer; +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; +import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.Float2Vector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.NullVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeNanoVector; +import org.apache.arrow.vector.TimeSecVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; +import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TimeStampNanoTZVector; +import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.arrow.vector.TimeStampSecTZVector; +import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.UInt1Vector; +import org.apache.arrow.vector.UInt2Vector; +import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.FixedSizeListVector; +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; +import org.apache.avro.SchemaBuilder; + +public class ArrowToAvroUtils { + + public static final String GENERIC_RECORD_TYPE_NAME = "GenericRecord"; + + /** + * Create an Avro record schema for a given list of Arrow fields. + * + *

This method currently performs following type mapping for Avro data types to corresponding + * Arrow data types. + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
Arrow typeAvro encoding
ArrowType.NullNULL
ArrowType.BoolBOOLEAN
ArrowType.Int(64 bit, unsigned 32 bit)LONG
ArrowType.Int(signed 32 bit, < 32 bit)INT
ArrowType.FloatingPoint(double)DOUBLE
ArrowType.FloatingPoint(single, half)FLOAT
ArrowType.Utf8STRING
ArrowType.LargeUtf8STRING
ArrowType.BinaryBYTES
ArrowType.LargeBinaryBYTES
ArrowType.FixedSizeBinaryFIXED
ArrowType.Decimaldecimal (FIXED)
ArrowType.Datedate (INT)
ArrowType.Time (SEC | MILLI)time-millis (INT)
ArrowType.Time (MICRO | NANO)time-micros (LONG)
ArrowType.Timestamp (NANOSECONDS, TZ != NULL)time-nanos (LONG)
ArrowType.Timestamp (MICROSECONDS, TZ != NULL)time-micros (LONG)
ArrowType.Timestamp (MILLISECONDS | SECONDS, TZ != NULL)time-millis (LONG)
ArrowType.Timestamp (NANOSECONDS, TZ == NULL)local-time-nanos (LONG)
ArrowType.Timestamp (MICROSECONDS, TZ == NULL)local-time-micros (LONG)
ArrowType.Timestamp (MILLISECONDS | SECONDS, TZ == NULL)local-time-millis (LONG)
ArrowType.Durationduration (FIXED)
ArrowType.Intervalduration (FIXED)
ArrowType.Structrecord
ArrowType.Listarray
ArrowType.LargeListarray
ArrowType.FixedSizeListarray
ArrowType.Mapmap
ArrowType.Unionunion
+ * + *

Nullable fields are represented as a union of [base-type | null]. Special treatment is given + * to nullability of unions - a union is considered nullable if any of its child fields are + * nullable. The schema for a nullable union will always contain a null type as its first member, + * with none of the child types being nullable. + * + *

List fields must contain precisely one child field, which may be nullable. Map fields are + * represented as a list of structs, where the struct fields are "key" and "value". The key field + * must always be of type STRING (Utf8) and cannot be nullable. The value can be of any type and + * 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, DictionaryProvider dictionaries) { + SchemaBuilder.RecordBuilder assembler = + SchemaBuilder.record(typeName).namespace(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. */ + public static Schema createAvroSchema(List arrowFields, String typeName) { + return createAvroSchema(arrowFields, typeName, null); + } + + /** Overload provided for convenience, sets name = GENERIC_RECORD_TYPE_NAME. */ + 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, + 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, dictionaries); + } + return assembler.endRecord(); + } + + private static SchemaBuilder.FieldAssembler buildFieldSchema( + SchemaBuilder.FieldAssembler assembler, + Field field, + String namespace, + DictionaryProvider dictionaries) { + + return assembler + .name(field.getName()) + .type(buildTypeSchema(SchemaBuilder.builder(), field, namespace, dictionaries)) + .noDefault(); + } + + private static T buildTypeSchema( + 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, 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, dictionaries); + return addTypesToUnion(union, tailTypes, namespace, dictionaries); + } + } else if (field.isNullable()) { + return buildBaseTypeSchema(builder.nullable(), field, namespace, dictionaries); + } else { + return buildBaseTypeSchema(builder, field, namespace, dictionaries); + } + } + + private static T buildArraySchema( + 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, dictionaries); + } + + private static T buildMapSchema( + 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"); + } + Field entriesField = mapField.getChildren().get(0); + if (mapField.getChildren().size() != 1) { + throw new IllegalArgumentException("Map entries must have exactly two child fields"); + } + Field keyField = entriesField.getChildren().get(0); + Field valueField = entriesField.getChildren().get(1); + if (keyField.getType().getTypeID() != ArrowType.ArrowTypeID.Utf8 || keyField.isNullable()) { + throw new IllegalArgumentException( + "Map keys must be of type string and cannot be nullable for conversion to Avro"); + } + return buildTypeSchema(builder.values(), valueField, namespace, dictionaries); + } + + private static T buildBaseTypeSchema( + SchemaBuilder.BaseTypeBuilder builder, + Field field, + String namespace, + DictionaryProvider dictionaries) { + + ArrowType.ArrowTypeID typeID = field.getType().getTypeID(); + + switch (typeID) { + case Null: + return builder.nullType(); + + case Bool: + 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(); + } else { + return builder.intType(); + } + + case FloatingPoint: + ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) field.getType(); + if (floatType.getPrecision() == FloatingPointPrecision.DOUBLE) { + return builder.doubleType(); + } else { + return builder.floatType(); + } + + case Utf8: + return builder.stringType(); + + case Binary: + return builder.bytesType(); + + case FixedSizeBinary: + ArrowType.FixedSizeBinary fixedType = (ArrowType.FixedSizeBinary) field.getType(); + String fixedTypeName = field.getName(); + int fixedTypeWidth = fixedType.getByteWidth(); + return builder.fixed(fixedTypeName).size(fixedTypeWidth); + + case Decimal: + ArrowType.Decimal decimalType = (ArrowType.Decimal) field.getType(); + return builder.type( + LogicalTypes.decimal(decimalType.getPrecision(), decimalType.getScale()) + .addToSchema( + Schema.createFixed( + field.getName(), namespace, "", decimalType.getBitWidth() / 8))); + + case Date: + return builder.type(LogicalTypes.date().addToSchema(Schema.create(Schema.Type.INT))); + + case Time: + ArrowType.Time timeType = (ArrowType.Time) field.getType(); + if ((timeType.getUnit() == TimeUnit.SECOND || timeType.getUnit() == TimeUnit.MILLISECOND)) { + // Second and millisecond time types are encoded as time-millis (INT) + return builder.type( + LogicalTypes.timeMillis().addToSchema(Schema.create(Schema.Type.INT))); + } else { + // All other time types (micro, nano) are encoded as time-micros (LONG) + return builder.type( + LogicalTypes.timeMicros().addToSchema(Schema.create(Schema.Type.LONG))); + } + + case Timestamp: + ArrowType.Timestamp timestampType = (ArrowType.Timestamp) field.getType(); + LogicalType timestampLogicalType = timestampLogicalType(timestampType); + return builder.type(timestampLogicalType.addToSchema(Schema.create(Schema.Type.LONG))); + + case Struct: + String childNamespace = + namespace == null ? field.getName() : namespace + "." + field.getName(); + return buildRecordSchema( + builder.record(field.getName()), field.getChildren(), childNamespace, dictionaries); + + case List: + case FixedSizeList: + // Arrow uses "$data$" as the field name for list items, that is not a valid Avro name + Field itemField = field.getChildren().get(0); + if (ListVector.DATA_VECTOR_NAME.equals(itemField.getName())) { + Field safeItemField = + new Field("item", itemField.getFieldType(), itemField.getChildren()); + Field safeListField = + new Field(field.getName(), field.getFieldType(), List.of(safeItemField)); + return buildArraySchema(builder.array(), safeListField, namespace, dictionaries); + } else { + return buildArraySchema(builder.array(), field, namespace, dictionaries); + } + + case Map: + return buildMapSchema(builder.map(), field, namespace, dictionaries); + + default: + throw new IllegalArgumentException( + "Element type not supported for Avro conversion: " + typeID.name()); + } + } + + private static T addTypesToUnion( + SchemaBuilder.UnionAccumulator accumulator, + List unionFields, + String namespace, + DictionaryProvider dictionaries) { + for (var field : unionFields) { + accumulator = buildBaseTypeSchema(accumulator.and(), field, namespace, dictionaries); + } + return accumulator.endUnion(); + } + + private static LogicalType timestampLogicalType(ArrowType.Timestamp timestampType) { + boolean zoneAware = timestampType.getTimezone() != null; + if (timestampType.getUnit() == TimeUnit.NANOSECOND) { + return zoneAware ? LogicalTypes.timestampNanos() : LogicalTypes.localTimestampNanos(); + } else if (timestampType.getUnit() == TimeUnit.MICROSECOND) { + return zoneAware ? LogicalTypes.timestampMicros() : LogicalTypes.localTimestampMicros(); + } else { + // Timestamp in seconds will be cast to milliseconds, Avro does not support seconds + return zoneAware ? LogicalTypes.timestampMillis() : LogicalTypes.localTimestampMillis(); + } + } + + 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, DictionaryProvider dictionaries) { + + List> producers = new ArrayList<>(vectors.size()); + + for (FieldVector vector : vectors) { + BaseAvroProducer producer = createProducer(vector, dictionaries); + producers.add(producer); + } + + return new CompositeAvroProducer(producers); + } + + /** 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, dictionaries); + } + + private static BaseAvroProducer createProducer( + FieldVector vector, boolean nullable, DictionaryProvider dictionaries) { + + Preconditions.checkNotNull(vector, "Arrow vector object can't be null"); + + final Types.MinorType minorType = vector.getMinorType(); + + // 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, 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); + case BIT: + return new AvroBooleanProducer((BitVector) vector); + case TINYINT: + return new AvroTinyIntProducer((TinyIntVector) vector); + case SMALLINT: + return new AvroSmallIntProducer((SmallIntVector) vector); + case INT: + return new AvroIntProducer((IntVector) vector); + case BIGINT: + return new AvroBigIntProducer((BigIntVector) vector); + case UINT1: + return new AvroUint1Producer((UInt1Vector) vector); + case UINT2: + return new AvroUint2Producer((UInt2Vector) vector); + case UINT4: + return new AvroUint4Producer((UInt4Vector) vector); + case UINT8: + return new AvroUint8Producer((UInt8Vector) vector); + case FLOAT2: + return new AvroFloat2Producer((Float2Vector) vector); + case FLOAT4: + return new AvroFloat4Producer((Float4Vector) vector); + case FLOAT8: + return new AvroFloat8Producer((Float8Vector) vector); + case VARBINARY: + return new AvroBytesProducer((VarBinaryVector) vector); + case FIXEDSIZEBINARY: + return new AvroFixedSizeBinaryProducer((FixedSizeBinaryVector) vector); + case VARCHAR: + return new AvroStringProducer((VarCharVector) vector); + + // Logical types + + case DECIMAL: + return new AvroDecimalProducer((DecimalVector) vector); + case DECIMAL256: + return new AvroDecimal256Producer((Decimal256Vector) vector); + case DATEDAY: + return new AvroDateDayProducer((DateDayVector) vector); + case DATEMILLI: + return new AvroDateMilliProducer((DateMilliVector) vector); + case TIMESEC: + return new AvroTimeSecProducer((TimeSecVector) vector); + case TIMEMILLI: + return new AvroTimeMilliProducer((TimeMilliVector) vector); + case TIMEMICRO: + return new AvroTimeMicroProducer((TimeMicroVector) vector); + case TIMENANO: + return new AvroTimeNanoProducer((TimeNanoVector) vector); + case TIMESTAMPSEC: + return new AvroTimestampSecProducer((TimeStampSecVector) vector); + case TIMESTAMPMILLI: + return new AvroTimestampMilliProducer((TimeStampMilliVector) vector); + case TIMESTAMPMICRO: + return new AvroTimestampMicroProducer((TimeStampMicroVector) vector); + case TIMESTAMPNANO: + return new AvroTimestampNanoProducer((TimeStampNanoVector) vector); + case TIMESTAMPSECTZ: + return new AvroTimestampSecTzProducer((TimeStampSecTZVector) vector); + case TIMESTAMPMILLITZ: + return new AvroTimestampMilliTzProducer((TimeStampMilliTZVector) vector); + case TIMESTAMPMICROTZ: + return new AvroTimestampMicroTzProducer((TimeStampMicroTZVector) vector); + case TIMESTAMPNANOTZ: + return new AvroTimestampNanoTzProducer((TimeStampNanoTZVector) vector); + + // Complex types + + case STRUCT: + StructVector structVector = (StructVector) vector; + List childVectors = structVector.getChildrenFromFields(); + 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(), dictionaries); + } + return new AvroStructProducer(structVector, childProducers); + + case LIST: + ListVector listVector = (ListVector) vector; + FieldVector itemVector = listVector.getDataVector(); + 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(), dictionaries); + return new AvroFixedSizeListProducer(fixedListVector, fixedItemProducer); + + case MAP: + MapVector mapVector = (MapVector) vector; + StructVector entryVector = (StructVector) mapVector.getDataVector(); + Types.MinorType keyType = entryVector.getChildrenFromFields().get(0).getMinorType(); + if (keyType != Types.MinorType.VARCHAR) { + throw new IllegalArgumentException("MAP key type must be VARCHAR for Avro encoding"); + } + VarCharVector keyVector = (VarCharVector) entryVector.getChildrenFromFields().get(0); + FieldVector valueVector = entryVector.getChildrenFromFields().get(1); + Producer keyProducer = new AvroStringProducer(keyVector); + Producer valueProducer = + createProducer(valueVector, valueVector.getField().isNullable(), dictionaries); + Producer entryProducer = + new AvroStructProducer(entryVector, new Producer[] {keyProducer, valueProducer}); + return new AvroMapProducer(mapVector, entryProducer); + + // Support for UNION and DENSEUNION is not currently available + // This is pending fixes in the implementation of the union vectors themselves + // https://github.com/apache/arrow-java/issues/108 + + default: + // Not all Arrow types are supported for encoding (yet)! + String error = + String.format( + "Encoding Arrow type %s to Avro is not currently supported", minorType.name()); + throw new UnsupportedOperationException(error); + } + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrow.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrow.java index 2392c36f94..2a28ad393b 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrow.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrow.java @@ -59,4 +59,23 @@ public static AvroToArrowVectorIterator avroToArrowIterator( return AvroToArrowVectorIterator.create(decoder, schema, config); } + + /** + * Convert an Avro schema to its Arrow equivalent. + * + *

The resulting set of Arrow fields matches what would be set in the VSR after calling + * avroToArrow() or avroToArrowIterator(), respecting the configuration in the config parameter. + * + * @param schema The Avro schema to convert + * @param config Configuration options for conversion + * @return The equivalent Arrow schema + */ + public static org.apache.arrow.vector.types.pojo.Schema avroToAvroSchema( + Schema schema, AvroToArrowConfig config) { + + Preconditions.checkNotNull(schema, "Avro schema object cannot be null"); + Preconditions.checkNotNull(config, "config cannot be null"); + + return AvroToArrowUtils.createArrowSchema(schema, config); + } } diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowConfig.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowConfig.java index 290d1a77d9..5596138586 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowConfig.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowConfig.java @@ -25,6 +25,7 @@ public class AvroToArrowConfig { private final BufferAllocator allocator; + /** * The maximum rowCount to read each time when partially convert data. Default value is 1024 and * -1 means read all data into one vector. @@ -40,6 +41,12 @@ public class AvroToArrowConfig { /** The field names which to skip when reading decoder values. */ private final Set skipFieldNames; + /** + * Use legacy-mode to keep compatibility with old behavior (pre-2025), enabled by default. This + * affects how the AvroToArrow code interprets the Avro schema. + */ + private final boolean legacyMode; + /** * Instantiate an instance. * @@ -63,6 +70,37 @@ public class AvroToArrowConfig { this.targetBatchSize = targetBatchSize; this.provider = provider; this.skipFieldNames = skipFieldNames; + + // Default values for optional parameters + legacyMode = true; // Keep compatibility with old behavior by default + } + + /** + * Instantiate an instance. + * + * @param allocator The memory allocator to construct the Arrow vectors with. + * @param targetBatchSize The maximum rowCount to read each time when partially convert data. + * @param provider The dictionary provider used for enum type, adapter will update this provider. + * @param skipFieldNames Field names which to skip. + * @param legacyMode Keep compatibility with old behavior (pre-2025) + */ + AvroToArrowConfig( + BufferAllocator allocator, + int targetBatchSize, + DictionaryProvider.MapDictionaryProvider provider, + Set skipFieldNames, + boolean legacyMode) { + + Preconditions.checkArgument( + targetBatchSize == AvroToArrowVectorIterator.NO_LIMIT_BATCH_SIZE || targetBatchSize > 0, + "invalid targetBatchSize: %s", + targetBatchSize); + + this.allocator = allocator; + this.targetBatchSize = targetBatchSize; + this.provider = provider; + this.skipFieldNames = skipFieldNames; + this.legacyMode = legacyMode; } public BufferAllocator getAllocator() { @@ -80,4 +118,8 @@ public DictionaryProvider.MapDictionaryProvider getProvider() { public Set getSkipFieldNames() { return skipFieldNames; } + + public boolean isLegacyMode() { + return legacyMode; + } } 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 b39121cfd1..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 @@ -41,6 +41,7 @@ import org.apache.arrow.adapter.avro.consumers.AvroLongConsumer; import org.apache.arrow.adapter.avro.consumers.AvroMapConsumer; import org.apache.arrow.adapter.avro.consumers.AvroNullConsumer; +import org.apache.arrow.adapter.avro.consumers.AvroNullableConsumer; import org.apache.arrow.adapter.avro.consumers.AvroStringConsumer; import org.apache.arrow.adapter.avro.consumers.AvroStructConsumer; import org.apache.arrow.adapter.avro.consumers.AvroUnionsConsumer; @@ -49,17 +50,23 @@ import org.apache.arrow.adapter.avro.consumers.SkipConsumer; import org.apache.arrow.adapter.avro.consumers.SkipFunction; import org.apache.arrow.adapter.avro.consumers.logical.AvroDateConsumer; +import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimal256Consumer; import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimalConsumer; import org.apache.arrow.adapter.avro.consumers.logical.AvroTimeMicroConsumer; import org.apache.arrow.adapter.avro.consumers.logical.AvroTimeMillisConsumer; import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampMicrosConsumer; +import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampMicrosTzConsumer; import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampMillisConsumer; +import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampMillisTzConsumer; +import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampNanosConsumer; +import org.apache.arrow.adapter.avro.consumers.logical.AvroTimestampNanosTzConsumer; import org.apache.arrow.memory.BufferAllocator; 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; +import org.apache.arrow.vector.Decimal256Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -69,8 +76,12 @@ import org.apache.arrow.vector.NullVector; import org.apache.arrow.vector.TimeMicroVector; import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TimeStampNanoTZVector; +import org.apache.arrow.vector.TimeStampNanoVector; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; @@ -169,42 +180,69 @@ private static Consumer createConsumer( switch (type) { case UNION: - consumer = createUnionConsumer(schema, name, config, consumerVector); + boolean nullableUnion = + schema.getTypes().stream().anyMatch(t -> t.getType() == Schema.Type.NULL); + if (schema.getTypes().size() == 2 && nullableUnion && !config.isLegacyMode()) { + // For a simple nullable (null | type), interpret the union as a single nullable field. + // Not available in legacy mode, which uses the literal interpretation instead + int nullIndex = schema.getTypes().get(0).getType() == Schema.Type.NULL ? 0 : 1; + int childIndex = nullIndex == 0 ? 1 : 0; + Schema childSchema = schema.getTypes().get(childIndex); + Consumer childConsumer = + createConsumer(childSchema, name, true, config, consumerVector); + consumer = new AvroNullableConsumer<>(childConsumer, nullIndex); + } else { + // Literal interpretation of a union, which may or may not include a null element. + consumer = createUnionConsumer(schema, name, nullableUnion, config, consumerVector); + } break; case ARRAY: - consumer = createArrayConsumer(schema, name, config, consumerVector); + consumer = createArrayConsumer(schema, name, nullable, config, consumerVector); break; case MAP: - consumer = createMapConsumer(schema, name, config, consumerVector); + consumer = createMapConsumer(schema, name, nullable, config, consumerVector); break; case RECORD: - consumer = createStructConsumer(schema, name, config, consumerVector); + consumer = createStructConsumer(schema, name, nullable, config, consumerVector); break; case ENUM: - consumer = createEnumConsumer(schema, name, config, consumerVector); + consumer = createEnumConsumer(schema, name, nullable, config, consumerVector); break; case STRING: arrowType = new ArrowType.Utf8(); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType(nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroStringConsumer((VarCharVector) vector); break; case FIXED: - Map extProps = createExternalProps(schema); + Map extProps = createExternalProps(schema, config); if (logicalType instanceof LogicalTypes.Decimal) { - arrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType); + arrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType, schema); fieldType = new FieldType( - nullable, arrowType, /*dictionary=*/ null, getMetaData(schema, extProps)); + nullable, + arrowType, + /* dictionary= */ null, + getMetaData(schema, extProps, config)); vector = createVector(consumerVector, fieldType, name, allocator); - consumer = - new AvroDecimalConsumer.FixedDecimalConsumer( - (DecimalVector) vector, schema.getFixedSize()); + if (schema.getFixedSize() <= 16) { + consumer = + new AvroDecimalConsumer.FixedDecimalConsumer( + (DecimalVector) vector, schema.getFixedSize()); + } else { + consumer = + new AvroDecimal256Consumer.FixedDecimal256Consumer( + (Decimal256Vector) vector, schema.getFixedSize()); + } } else { arrowType = new ArrowType.FixedSizeBinary(schema.getFixedSize()); fieldType = new FieldType( - nullable, arrowType, /*dictionary=*/ null, getMetaData(schema, extProps)); + nullable, + arrowType, + /* dictionary= */ null, + getMetaData(schema, extProps, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroFixedConsumer((FixedSizeBinaryVector) vector, schema.getFixedSize()); } @@ -212,79 +250,141 @@ private static Consumer createConsumer( case INT: if (logicalType instanceof LogicalTypes.Date) { arrowType = new ArrowType.Date(DateUnit.DAY); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroDateConsumer((DateDayVector) vector); } else if (logicalType instanceof LogicalTypes.TimeMillis) { arrowType = new ArrowType.Time(TimeUnit.MILLISECOND, 32); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimeMillisConsumer((TimeMilliVector) vector); } else { - arrowType = new ArrowType.Int(32, /*isSigned=*/ true); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + arrowType = new ArrowType.Int(32, /* isSigned= */ true); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroIntConsumer((IntVector) vector); } break; case BOOLEAN: arrowType = new ArrowType.Bool(); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType(nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroBooleanConsumer((BitVector) vector); break; case LONG: if (logicalType instanceof LogicalTypes.TimeMicros) { arrowType = new ArrowType.Time(TimeUnit.MICROSECOND, 64); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimeMicroConsumer((TimeMicroVector) vector); - } else if (logicalType instanceof LogicalTypes.TimestampMillis) { + } else if (logicalType instanceof LogicalTypes.TimestampMillis && !config.isLegacyMode()) { + // In legacy mode the timestamp-xxx types are treated as local, there is no zone aware + // type + arrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); + vector = createVector(consumerVector, fieldType, name, allocator); + consumer = new AvroTimestampMillisTzConsumer((TimeStampMilliTZVector) vector); + } else if (logicalType instanceof LogicalTypes.TimestampMicros && !config.isLegacyMode()) { + arrowType = new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); + vector = createVector(consumerVector, fieldType, name, allocator); + consumer = new AvroTimestampMicrosTzConsumer((TimeStampMicroTZVector) vector); + } else if (logicalType instanceof LogicalTypes.TimestampNanos && !config.isLegacyMode()) { + arrowType = new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); + vector = createVector(consumerVector, fieldType, name, allocator); + consumer = new AvroTimestampNanosTzConsumer((TimeStampNanoTZVector) vector); + } else if (logicalType instanceof LogicalTypes.LocalTimestampMillis + || (logicalType instanceof LogicalTypes.TimestampMillis && config.isLegacyMode())) { arrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimestampMillisConsumer((TimeStampMilliVector) vector); - } else if (logicalType instanceof LogicalTypes.TimestampMicros) { + } else if (logicalType instanceof LogicalTypes.LocalTimestampMicros + || (logicalType instanceof LogicalTypes.TimestampMicros && config.isLegacyMode())) { + // In legacy mode the timestamp-xxx types are treated as local arrowType = new ArrowType.Timestamp(TimeUnit.MICROSECOND, null); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroTimestampMicrosConsumer((TimeStampMicroVector) vector); + } else if (logicalType instanceof LogicalTypes.LocalTimestampNanos + || (logicalType instanceof LogicalTypes.TimestampNanos && config.isLegacyMode())) { + arrowType = new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); + vector = createVector(consumerVector, fieldType, name, allocator); + consumer = new AvroTimestampNanosConsumer((TimeStampNanoVector) vector); } else { - arrowType = new ArrowType.Int(64, /*isSigned=*/ true); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + arrowType = new ArrowType.Int(64, /* isSigned= */ true); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroLongConsumer((BigIntVector) vector); } break; case FLOAT: arrowType = new ArrowType.FloatingPoint(SINGLE); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType(nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroFloatConsumer((Float4Vector) vector); break; case DOUBLE: arrowType = new ArrowType.FloatingPoint(DOUBLE); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType(nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroDoubleConsumer((Float8Vector) vector); break; case BYTES: if (logicalType instanceof LogicalTypes.Decimal) { - arrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + LogicalTypes.Decimal decimalType = (LogicalTypes.Decimal) logicalType; + arrowType = createDecimalArrowType(decimalType, schema); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); - consumer = new AvroDecimalConsumer.BytesDecimalConsumer((DecimalVector) vector); + if (decimalType.getPrecision() <= 38) { + consumer = new AvroDecimalConsumer.BytesDecimalConsumer((DecimalVector) vector); + } else { + consumer = + new AvroDecimal256Consumer.BytesDecimal256Consumer((Decimal256Vector) vector); + } } else { arrowType = new ArrowType.Binary(); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); + fieldType = + new FieldType( + nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); vector = createVector(consumerVector, fieldType, name, allocator); consumer = new AvroBytesConsumer((VarBinaryVector) vector); } break; case NULL: arrowType = new ArrowType.Null(); - fieldType = new FieldType(nullable, arrowType, /*dictionary=*/ null, getMetaData(schema)); - vector = fieldType.createNewSingleVector(name, allocator, /*schemaCallBack=*/ null); + fieldType = + new FieldType(nullable, arrowType, /* dictionary= */ null, getMetaData(schema, config)); + vector = new NullVector(name, fieldType); // Respect nullability defined in fieldType consumer = new AvroNullConsumer((NullVector) vector); break; default: @@ -295,19 +395,31 @@ private static Consumer createConsumer( return consumer; } - private static ArrowType createDecimalArrowType(LogicalTypes.Decimal logicalType) { + private static ArrowType createDecimalArrowType(LogicalTypes.Decimal logicalType, Schema schema) { final int scale = logicalType.getScale(); final int precision = logicalType.getPrecision(); Preconditions.checkArgument( - precision > 0 && precision <= 38, "Precision must be in range of 1 to 38"); - Preconditions.checkArgument(scale >= 0 && scale <= 38, "Scale must be in range of 0 to 38."); + precision > 0 && precision <= 76, "Precision must be in range of 1 to 76"); + Preconditions.checkArgument(scale >= 0 && scale <= 76, "Scale must be in range of 0 to 76."); Preconditions.checkArgument( scale <= precision, "Invalid decimal scale: %s (greater than precision: %s)", scale, precision); - return new ArrowType.Decimal(precision, scale, 128); + if (schema.getType() == Schema.Type.FIXED) { + if (schema.getFixedSize() <= 16) { + return new ArrowType.Decimal(precision, scale, 128); + } else { + return new ArrowType.Decimal(precision, scale, 256); + } + } else { + if (precision <= 38) { + return new ArrowType.Decimal(precision, scale, 128); + } else { + return new ArrowType.Decimal(precision, scale, 256); + } + } } private static Consumer createSkipConsumer(Schema schema) { @@ -397,6 +509,30 @@ private static Consumer createSkipConsumer(Schema schema) { return new SkipConsumer(skipFunction); } + static org.apache.arrow.vector.types.pojo.Schema createArrowSchema( + Schema schema, AvroToArrowConfig config) { + + // Create an Arrow schema matching the structure of vectors built by createCompositeConsumer() + + Set skipFieldNames = config.getSkipFieldNames(); + List arrowFields = new ArrayList<>(schema.getFields().size()); + + Schema.Type type = schema.getType(); + if (type == Schema.Type.RECORD) { + for (Schema.Field field : schema.getFields()) { + if (!skipFieldNames.contains(field.name())) { + Field arrowField = avroSchemaToField(field.schema(), field.name(), config); + arrowFields.add(arrowField); + } + } + } else { + Field arrowField = avroSchemaToField(schema, schema.getName(), config); + arrowFields.add(arrowField); + } + + return new org.apache.arrow.vector.types.pojo.Schema(arrowFields); + } + static CompositeAvroConsumer createCompositeConsumer(Schema schema, AvroToArrowConfig config) { List consumers = new ArrayList<>(); @@ -433,11 +569,20 @@ private static String getDefaultFieldName(ArrowType type) { } private static Field avroSchemaToField(Schema schema, String name, AvroToArrowConfig config) { - return avroSchemaToField(schema, name, config, null); + return avroSchemaToField(schema, name, false, config, null); } private static Field avroSchemaToField( Schema schema, String name, AvroToArrowConfig config, Map externalProps) { + return avroSchemaToField(schema, name, false, config, externalProps); + } + + private static Field avroSchemaToField( + Schema schema, + String name, + boolean nullable, + AvroToArrowConfig config, + Map externalProps) { final Schema.Type type = schema.getType(); final LogicalType logicalType = schema.getLogicalType(); @@ -446,33 +591,53 @@ private static Field avroSchemaToField( switch (type) { case UNION: - for (int i = 0; i < schema.getTypes().size(); i++) { - Schema childSchema = schema.getTypes().get(i); - // Union child vector should use default name - children.add(avroSchemaToField(childSchema, null, config)); + boolean nullableUnion = + schema.getTypes().stream().anyMatch(t -> t.getType() == Schema.Type.NULL); + if (nullableUnion && schema.getTypes().size() == 2 && !config.isLegacyMode()) { + // For a simple nullable (null | type), interpret the union as a single nullable field. + // Not available in legacy mode, which uses the literal interpretation instead + Schema childSchema = + schema.getTypes().get(0).getType() == Schema.Type.NULL + ? schema.getTypes().get(1) + : schema.getTypes().get(0); + return avroSchemaToField(childSchema, name, true, config, externalProps); + } else { + // Literal interpretation of a union, which may or may not include a null element. + for (int i = 0; i < schema.getTypes().size(); i++) { + Schema childSchema = schema.getTypes().get(i); + // Union child vector should use default name + children.add(avroSchemaToField(childSchema, null, nullableUnion, config, null)); + } + fieldType = + createFieldType( + new ArrowType.Union(UnionMode.Sparse, null), schema, externalProps, config); } - fieldType = - createFieldType(new ArrowType.Union(UnionMode.Sparse, null), schema, externalProps); break; case ARRAY: Schema elementSchema = schema.getElementType(); - children.add(avroSchemaToField(elementSchema, elementSchema.getName(), config)); - fieldType = createFieldType(new ArrowType.List(), schema, externalProps); + children.add(avroSchemaToField(elementSchema, ListVector.DATA_VECTOR_NAME, config)); + fieldType = createFieldType(nullable, new ArrowType.List(), schema, externalProps, config); break; case MAP: // MapVector internal struct field and key field should be non-nullable FieldType keyFieldType = - new FieldType(/*nullable=*/ false, new ArrowType.Utf8(), /*dictionary=*/ null); - Field keyField = new Field("key", keyFieldType, /*children=*/ null); - Field valueField = avroSchemaToField(schema.getValueType(), "value", config); + new FieldType(/* nullable= */ false, new ArrowType.Utf8(), /* dictionary= */ null); + Field keyField = new Field(MapVector.KEY_NAME, keyFieldType, /* children= */ null); + Field valueField = avroSchemaToField(schema.getValueType(), MapVector.VALUE_NAME, config); FieldType structFieldType = - new FieldType(false, new ArrowType.Struct(), /*dictionary=*/ null); + new FieldType(false, new ArrowType.Struct(), /* dictionary= */ null); Field structField = - new Field("internal", structFieldType, Arrays.asList(keyField, valueField)); + new Field( + MapVector.DATA_VECTOR_NAME, structFieldType, Arrays.asList(keyField, valueField)); children.add(structField); fieldType = - createFieldType(new ArrowType.Map(/*keysSorted=*/ false), schema, externalProps); + createFieldType( + nullable, + new ArrowType.Map(/* keysSorted= */ false), + schema, + externalProps, + config); break; case RECORD: final Set skipFieldNames = config.getSkipFieldNames(); @@ -487,13 +652,14 @@ private static Field avroSchemaToField( if (doc != null) { extProps.put("doc", doc); } - if (aliases != null) { + if (aliases != null && (!aliases.isEmpty() || config.isLegacyMode())) { extProps.put("aliases", convertAliases(aliases)); } children.add(avroSchemaToField(childSchema, fullChildName, config, extProps)); } } - fieldType = createFieldType(new ArrowType.Struct(), schema, externalProps); + fieldType = + createFieldType(nullable, new ArrowType.Struct(), schema, externalProps, config); break; case ENUM: DictionaryProvider.MapDictionaryProvider provider = config.getProvider(); @@ -503,23 +669,25 @@ private static Field avroSchemaToField( fieldType = createFieldType( + nullable, indexType, schema, externalProps, - new DictionaryEncoding(current, /*ordered=*/ false, /*indexType=*/ indexType)); + new DictionaryEncoding(current, /* ordered= */ false, /* indexType= */ indexType), + config); break; case STRING: - fieldType = createFieldType(new ArrowType.Utf8(), schema, externalProps); + fieldType = createFieldType(nullable, new ArrowType.Utf8(), schema, externalProps, config); break; case FIXED: final ArrowType fixedArrowType; if (logicalType instanceof LogicalTypes.Decimal) { - fixedArrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType); + fixedArrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType, schema); } else { fixedArrowType = new ArrowType.FixedSizeBinary(schema.getFixedSize()); } - fieldType = createFieldType(fixedArrowType, schema, externalProps); + fieldType = createFieldType(nullable, fixedArrowType, schema, externalProps, config); break; case INT: final ArrowType intArrowType; @@ -528,43 +696,64 @@ private static Field avroSchemaToField( } else if (logicalType instanceof LogicalTypes.TimeMillis) { intArrowType = new ArrowType.Time(TimeUnit.MILLISECOND, 32); } else { - intArrowType = new ArrowType.Int(32, /*isSigned=*/ true); + intArrowType = new ArrowType.Int(32, /* isSigned= */ true); } - fieldType = createFieldType(intArrowType, schema, externalProps); + fieldType = createFieldType(nullable, intArrowType, schema, externalProps, config); break; case BOOLEAN: - fieldType = createFieldType(new ArrowType.Bool(), schema, externalProps); + fieldType = createFieldType(nullable, new ArrowType.Bool(), schema, externalProps, config); break; case LONG: final ArrowType longArrowType; if (logicalType instanceof LogicalTypes.TimeMicros) { longArrowType = new ArrowType.Time(TimeUnit.MICROSECOND, 64); } else if (logicalType instanceof LogicalTypes.TimestampMillis) { - longArrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + // In legacy mode the timestamp-xxx types are treated as local + String tz = config.isLegacyMode() ? null : "UTC"; + longArrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, tz); } else if (logicalType instanceof LogicalTypes.TimestampMicros) { + String tz = config.isLegacyMode() ? null : "UTC"; + longArrowType = new ArrowType.Timestamp(TimeUnit.MICROSECOND, tz); + } else if (logicalType instanceof LogicalTypes.TimestampNanos) { + String tz = config.isLegacyMode() ? null : "UTC"; + longArrowType = new ArrowType.Timestamp(TimeUnit.NANOSECOND, tz); + } else if (logicalType instanceof LogicalTypes.LocalTimestampMillis + && !config.isLegacyMode()) { + // In legacy mode the local-timestamp-xxx types are not recognized (result is just type = + // long) + longArrowType = new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + } else if (logicalType instanceof LogicalTypes.LocalTimestampMicros + && !config.isLegacyMode()) { longArrowType = new ArrowType.Timestamp(TimeUnit.MICROSECOND, null); + } else if (logicalType instanceof LogicalTypes.LocalTimestampNanos + && !config.isLegacyMode()) { + longArrowType = new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); } else { - longArrowType = new ArrowType.Int(64, /*isSigned=*/ true); + longArrowType = new ArrowType.Int(64, /* isSigned= */ true); } - fieldType = createFieldType(longArrowType, schema, externalProps); + fieldType = createFieldType(nullable, longArrowType, schema, externalProps, config); break; case FLOAT: - fieldType = createFieldType(new ArrowType.FloatingPoint(SINGLE), schema, externalProps); + fieldType = + createFieldType( + nullable, new ArrowType.FloatingPoint(SINGLE), schema, externalProps, config); break; case DOUBLE: - fieldType = createFieldType(new ArrowType.FloatingPoint(DOUBLE), schema, externalProps); + fieldType = + createFieldType( + nullable, new ArrowType.FloatingPoint(DOUBLE), schema, externalProps, config); break; case BYTES: final ArrowType bytesArrowType; if (logicalType instanceof LogicalTypes.Decimal) { - bytesArrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType); + bytesArrowType = createDecimalArrowType((LogicalTypes.Decimal) logicalType, schema); } else { bytesArrowType = new ArrowType.Binary(); } - fieldType = createFieldType(bytesArrowType, schema, externalProps); + fieldType = createFieldType(nullable, bytesArrowType, schema, externalProps, config); break; case NULL: - fieldType = createFieldType(ArrowType.Null.INSTANCE, schema, externalProps); + fieldType = createFieldType(ArrowType.Null.INSTANCE, schema, externalProps, config); break; default: // no-op, shouldn't get here @@ -574,15 +763,24 @@ private static Field avroSchemaToField( if (name == null) { name = getDefaultFieldName(fieldType.getType()); } + if (name.contains(".") && !config.isLegacyMode()) { + // Do not include namespace as part of the field name + name = name.substring(name.lastIndexOf(".") + 1); + } return new Field(name, fieldType, children.size() == 0 ? null : children); } private static Consumer createArrayConsumer( - Schema schema, String name, AvroToArrowConfig config, FieldVector consumerVector) { + Schema schema, + String name, + boolean nullable, + AvroToArrowConfig config, + FieldVector consumerVector) { ListVector listVector; if (consumerVector == null) { - final Field field = avroSchemaToField(schema, name, config); + final Field field = + avroSchemaToField(schema, name, nullable, config, /* externalProps= */ null); listVector = (ListVector) field.createVector(config.getAllocator()); } else { listVector = (ListVector) consumerVector; @@ -598,13 +796,18 @@ private static Consumer createArrayConsumer( } private static Consumer createStructConsumer( - Schema schema, String name, AvroToArrowConfig config, FieldVector consumerVector) { + Schema schema, + String name, + boolean nullable, + AvroToArrowConfig config, + FieldVector consumerVector) { final Set skipFieldNames = config.getSkipFieldNames(); StructVector structVector; if (consumerVector == null) { - final Field field = avroSchemaToField(schema, name, config, createExternalProps(schema)); + final Field field = + avroSchemaToField(schema, name, nullable, config, createExternalProps(schema, config)); structVector = (StructVector) field.createVector(config.getAllocator()); } else { structVector = (StructVector) consumerVector; @@ -635,11 +838,16 @@ private static Consumer createStructConsumer( } private static Consumer createEnumConsumer( - Schema schema, String name, AvroToArrowConfig config, FieldVector consumerVector) { + Schema schema, + String name, + boolean nullable, + AvroToArrowConfig config, + FieldVector consumerVector) { BaseIntVector indexVector; if (consumerVector == null) { - final Field field = avroSchemaToField(schema, name, config, createExternalProps(schema)); + final Field field = + avroSchemaToField(schema, name, nullable, config, createExternalProps(schema, config)); indexVector = (BaseIntVector) field.createVector(config.getAllocator()); } else { indexVector = (BaseIntVector) consumerVector; @@ -659,11 +867,16 @@ private static Consumer createEnumConsumer( } private static Consumer createMapConsumer( - Schema schema, String name, AvroToArrowConfig config, FieldVector consumerVector) { + Schema schema, + String name, + boolean nullable, + AvroToArrowConfig config, + FieldVector consumerVector) { MapVector mapVector; if (consumerVector == null) { - final Field field = avroSchemaToField(schema, name, config); + final Field field = + avroSchemaToField(schema, name, nullable, config, /* externalProps= */ null); mapVector = (MapVector) field.createVector(config.getAllocator()); } else { mapVector = (MapVector) consumerVector; @@ -689,12 +902,13 @@ private static Consumer createMapConsumer( } private static Consumer createUnionConsumer( - Schema schema, String name, AvroToArrowConfig config, FieldVector consumerVector) { + Schema schema, + String name, + boolean nullableUnion, + AvroToArrowConfig config, + FieldVector consumerVector) { final int size = schema.getTypes().size(); - final boolean nullable = - schema.getTypes().stream().anyMatch(t -> t.getType() == Schema.Type.NULL); - UnionVector unionVector; if (consumerVector == null) { final Field field = avroSchemaToField(schema, name, config); @@ -711,7 +925,8 @@ private static Consumer createUnionConsumer( for (int i = 0; i < size; i++) { FieldVector child = childVectors.get(i); Schema subSchema = schema.getTypes().get(i); - Consumer delegate = createConsumer(subSchema, subSchema.getName(), nullable, config, child); + Consumer delegate = + createConsumer(subSchema, subSchema.getName(), nullableUnion, config, child); delegates[i] = delegate; types[i] = child.getMinorType(); } @@ -776,14 +991,24 @@ static VectorSchemaRoot avroToArrowVectors( return root; } - private static Map getMetaData(Schema schema) { + // Do not include props that are part of the Avro format itself as field metadata + // These are already represented in the field / type structure and are not custom attributes + private static final List AVRO_FORMAT_METADATA = + Arrays.asList("logicalType", "precision", "scale"); + + private static Map getMetaData(Schema schema, AvroToArrowConfig config) { Map metadata = new HashMap<>(); - schema.getObjectProps().forEach((k, v) -> metadata.put(k, v.toString())); + for (Map.Entry prop : schema.getObjectProps().entrySet()) { + if (!AVRO_FORMAT_METADATA.contains(prop.getKey()) || config.isLegacyMode()) { + metadata.put(prop.getKey(), prop.getValue().toString()); + } + } return metadata; } - private static Map getMetaData(Schema schema, Map externalProps) { - Map metadata = getMetaData(schema); + private static Map getMetaData( + Schema schema, Map externalProps, AvroToArrowConfig config) { + Map metadata = getMetaData(schema, config); if (externalProps != null) { metadata.putAll(externalProps); } @@ -791,37 +1016,63 @@ private static Map getMetaData(Schema schema, Map createExternalProps(Schema schema) { + private static Map createExternalProps(Schema schema, AvroToArrowConfig config) { final Map extProps = new HashMap<>(); String doc = schema.getDoc(); Set aliases = schema.getAliases(); if (doc != null) { extProps.put("doc", doc); } - if (aliases != null) { + if (aliases != null && (!aliases.isEmpty() || config.isLegacyMode())) { extProps.put("aliases", convertAliases(aliases)); } return extProps; } private static FieldType createFieldType( - ArrowType arrowType, Schema schema, Map externalProps) { - return createFieldType(arrowType, schema, externalProps, /*dictionary=*/ null); + ArrowType arrowType, + Schema schema, + Map externalProps, + AvroToArrowConfig config) { + return createFieldType(arrowType, schema, externalProps, /* dictionary= */ null, config); + } + + private static FieldType createFieldType( + boolean nullable, + ArrowType arrowType, + Schema schema, + Map externalProps, + AvroToArrowConfig config) { + return createFieldType( + nullable, arrowType, schema, externalProps, /* dictionary= */ null, config); + } + + private static FieldType createFieldType( + ArrowType arrowType, + Schema schema, + Map externalProps, + DictionaryEncoding dictionary, + AvroToArrowConfig config) { + + return createFieldType( + /* nullable= */ false, arrowType, schema, externalProps, dictionary, config); } private static FieldType createFieldType( + boolean nullable, ArrowType arrowType, Schema schema, Map externalProps, - DictionaryEncoding dictionary) { + DictionaryEncoding dictionary, + AvroToArrowConfig config) { return new FieldType( - /*nullable=*/ false, arrowType, dictionary, getMetaData(schema, externalProps)); + nullable, arrowType, dictionary, getMetaData(schema, externalProps, config)); } 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/consumers/AvroNullableConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/AvroNullableConsumer.java new file mode 100644 index 0000000000..b67819cb9d --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/AvroNullableConsumer.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.consumers; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; +import org.apache.avro.io.Decoder; + +/** + * Consumer wrapper which consumes nullable type values from avro decoder. Write the data to the + * underlying {@link FieldVector}. + * + * @param The vector within consumer or its delegate. + */ +public class AvroNullableConsumer extends BaseAvroConsumer { + + private final Consumer delegate; + private final int nullIndex; + + /** Instantiate a AvroNullableConsumer. */ + @SuppressWarnings("unchecked") + public AvroNullableConsumer(Consumer delegate, int nullIndex) { + super((T) delegate.getVector()); + this.delegate = delegate; + this.nullIndex = nullIndex; + } + + @Override + public void consume(Decoder decoder) throws IOException { + int typeIndex = decoder.readInt(); + if (typeIndex == nullIndex) { + decoder.readNull(); + delegate.addNull(); + } else { + delegate.consume(decoder); + } + currentIndex++; + } + + @Override + public void addNull() { + // Can be called by containers of nullable types + delegate.addNull(); + currentIndex++; + } + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + delegate.setPosition(index); + super.setPosition(index); + } + + @Override + public boolean resetValueVector(T vector) { + boolean delegateOk = delegate.resetValueVector(vector); + boolean thisOk = super.resetValueVector(vector); + return thisOk && delegateOk; + } + + @Override + public void close() throws Exception { + super.close(); + delegate.close(); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroDecimal256Consumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroDecimal256Consumer.java new file mode 100644 index 0000000000..12652833a1 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroDecimal256Consumer.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.consumers.logical; + +import java.io.IOException; +import java.nio.ByteBuffer; +import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer; +import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.Decimal256Vector; +import org.apache.avro.io.Decoder; + +/** + * Consumer which consume 256-bit decimal type values from avro decoder. Write the data to {@link + * Decimal256Vector}. + */ +public abstract class AvroDecimal256Consumer extends BaseAvroConsumer { + + protected AvroDecimal256Consumer(Decimal256Vector vector) { + super(vector); + } + + /** Consumer for decimal logical type with 256 bit width and original bytes type. */ + public static class BytesDecimal256Consumer extends AvroDecimal256Consumer { + + private ByteBuffer cacheBuffer; + + /** Instantiate a BytesDecimal256Consumer. */ + public BytesDecimal256Consumer(Decimal256Vector vector) { + super(vector); + } + + @Override + public void consume(Decoder decoder) throws IOException { + cacheBuffer = decoder.readBytes(cacheBuffer); + byte[] bytes = new byte[cacheBuffer.limit()]; + Preconditions.checkArgument(bytes.length <= 32, "Decimal bytes length should <= 32."); + cacheBuffer.get(bytes); + vector.setBigEndian(currentIndex++, bytes); + } + } + + /** Consumer for decimal logical type with 256 bit width and original fixed type. */ + public static class FixedDecimal256Consumer extends AvroDecimal256Consumer { + + private final byte[] reuseBytes; + + /** Instantiate a FixedDecimal256Consumer. */ + public FixedDecimal256Consumer(Decimal256Vector vector, int size) { + super(vector); + Preconditions.checkArgument(size <= 32, "Decimal bytes length should <= 32."); + reuseBytes = new byte[size]; + } + + @Override + public void consume(Decoder decoder) throws IOException { + decoder.readFixed(reuseBytes); + vector.setBigEndian(currentIndex++, reuseBytes); + } + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosConsumer.java index 88acf7b329..5af40ed17d 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosConsumer.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosConsumer.java @@ -22,7 +22,7 @@ import org.apache.avro.io.Decoder; /** - * Consumer which consume date timestamp-micro values from avro decoder. Write the data to {@link + * Consumer which consumes local-timestamp-micros values from avro decoder. Write the data to {@link * TimeStampMicroVector}. */ public class AvroTimestampMicrosConsumer extends BaseAvroConsumer { diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosTzConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosTzConsumer.java new file mode 100644 index 0000000000..a5dede4988 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMicrosTzConsumer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.consumers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.avro.io.Decoder; + +/** + * Consumer which consumes timestamp-micros values from avro decoder. Write the data to {@link + * TimeStampMicroTZVector}. + */ +public class AvroTimestampMicrosTzConsumer extends BaseAvroConsumer { + + /** Instantiate a AvroTimestampMicrosTzConsumer. */ + public AvroTimestampMicrosTzConsumer(TimeStampMicroTZVector vector) { + super(vector); + } + + @Override + public void consume(Decoder decoder) throws IOException { + vector.set(currentIndex++, decoder.readLong()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisConsumer.java index ec50d79023..bc451bd1dc 100644 --- a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisConsumer.java +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisConsumer.java @@ -22,7 +22,7 @@ import org.apache.avro.io.Decoder; /** - * Consumer which consume date timestamp-millis values from avro decoder. Write the data to {@link + * Consumer which consume local-timestamp-millis values from avro decoder. Write the data to {@link * TimeStampMilliVector}. */ public class AvroTimestampMillisConsumer extends BaseAvroConsumer { diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisTzConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisTzConsumer.java new file mode 100644 index 0000000000..255fe501fb --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampMillisTzConsumer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.consumers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer; +import org.apache.arrow.vector.TimeStampMilliTZVector; +import org.apache.avro.io.Decoder; + +/** + * Consumer which consume timestamp-millis values from avro decoder. Write the data to {@link + * TimeStampMilliTZVector}. + */ +public class AvroTimestampMillisTzConsumer extends BaseAvroConsumer { + + /** Instantiate a AvroTimestampMillisTzConsumer. */ + public AvroTimestampMillisTzConsumer(TimeStampMilliTZVector vector) { + super(vector); + } + + @Override + public void consume(Decoder decoder) throws IOException { + vector.set(currentIndex++, decoder.readLong()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosConsumer.java new file mode 100644 index 0000000000..b5044d221f --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosConsumer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.consumers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer; +import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.avro.io.Decoder; + +/** + * Consumer which consume local-timestamp-nanos values from avro decoder. Write the data to {@link + * TimeStampNanoVector}. + */ +public class AvroTimestampNanosConsumer extends BaseAvroConsumer { + + /** Instantiate a AvroTimestampNanosConsumer. */ + public AvroTimestampNanosConsumer(TimeStampNanoVector vector) { + super(vector); + } + + @Override + public void consume(Decoder decoder) throws IOException { + vector.set(currentIndex++, decoder.readLong()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosTzConsumer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosTzConsumer.java new file mode 100644 index 0000000000..3f42b7ccbb --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroTimestampNanosTzConsumer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.consumers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer; +import org.apache.arrow.vector.TimeStampNanoTZVector; +import org.apache.avro.io.Decoder; + +/** + * Consumer which consume timestamp-nanos values from avro decoder. Write the data to {@link + * TimeStampNanoTZVector}. + */ +public class AvroTimestampNanosTzConsumer extends BaseAvroConsumer { + + /** Instantiate a AvroTimestampNanosConsumer. */ + public AvroTimestampNanosTzConsumer(TimeStampNanoTZVector vector) { + super(vector); + } + + @Override + public void consume(Decoder decoder) throws IOException { + vector.set(currentIndex++, decoder.readLong()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBigIntProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBigIntProducer.java new file mode 100644 index 0000000000..9712e157c6 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBigIntProducer.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.BaseFixedWidthVector; +import org.apache.arrow.vector.BigIntVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces long values from a {@link BigIntVector}, writes data to an Avro encoder. + * + *

Logical types are also supported, for vectors derived from {@link BaseFixedWidthVector} where + * the internal representation matches BigIntVector and requires no conversion. + */ +public class AvroBigIntProducer extends BaseAvroProducer { + + /** Instantiate an AvroBigIntProducer. */ + public AvroBigIntProducer(BigIntVector vector) { + super(vector); + } + + /** Protected constructor for logical types with a long representation. */ + protected AvroBigIntProducer(BaseFixedWidthVector vector) { + super(vector); + if (vector.getTypeWidth() != BigIntVector.TYPE_WIDTH) { + throw new IllegalArgumentException( + "AvroBigIntProducer requires type width = " + BigIntVector.TYPE_WIDTH); + } + } + + @Override + public void produce(Encoder encoder) throws IOException { + long value = vector.getDataBuffer().getLong(currentIndex * (long) BigIntVector.TYPE_WIDTH); + encoder.writeLong(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBooleanProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBooleanProducer.java new file mode 100644 index 0000000000..523ddf110d --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBooleanProducer.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.BitVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces boolean values from a {@link BitVector}, writes data to an Avro encoder. + */ +public class AvroBooleanProducer extends BaseAvroProducer { + + /** Instantiate am AvroBooleanProducer. */ + public AvroBooleanProducer(BitVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + int bitValue = vector.get(currentIndex++); + encoder.writeBoolean(bitValue != 0); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBytesProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBytesProducer.java new file mode 100644 index 0000000000..e1fe6dddc6 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroBytesProducer.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.adapter.avro.producers; + +import java.io.IOException; +import java.nio.ByteBuffer; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces byte array values from a {@link VarBinaryVector}, writes data to an Avro + * encoder. + */ +public class AvroBytesProducer extends BaseAvroProducer { + + /** Instantiate an AvroBytesProducer. */ + public AvroBytesProducer(VarBinaryVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + // The nio ByteBuffer is created once per call, but underlying data is not copied + long offset = vector.getStartOffset(currentIndex); + long endOffset = vector.getEndOffset(currentIndex); + int length = (int) (endOffset - offset); + ByteBuffer nioBuffer = vector.getDataBuffer().nioBuffer(offset, length); + encoder.writeBytes(nioBuffer); + currentIndex++; + } +} 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 new file mode 100644 index 0000000000..eebfb7d241 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroEnumProducer.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.avro.io.Encoder; + +/** + * Producer that produces enum values from a dictionary-encoded {@link BaseIntVector}, writes data + * to an Avro encoder. + */ +public class AvroEnumProducer extends BaseAvroProducer { + + /** Instantiate an AvroEnumProducer. */ + public AvroEnumProducer(BaseIntVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + encoder.writeEnum((int) vector.getValueAsLong(currentIndex++)); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeBinaryProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeBinaryProducer.java new file mode 100644 index 0000000000..9fb877cfa0 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeBinaryProducer.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.BaseFixedWidthVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces fixed-size binary values from a {@link FixedSizeBinaryVector}, writes data + * to an Avro encoder. + * + *

Logical types are also supported, for vectors derived from {@link BaseFixedWidthVector} where + * the internal representation is fixed width bytes and requires no conversion. + */ +public class AvroFixedSizeBinaryProducer extends BaseAvroProducer { + + private final byte[] reuseBytes; + + /** Instantiate an AvroFixedSizeBinaryProducer. */ + public AvroFixedSizeBinaryProducer(FixedSizeBinaryVector vector) { + super(vector); + reuseBytes = new byte[vector.getTypeWidth()]; + } + + /** Protected constructor for logical types with a fixed width representation. */ + protected AvroFixedSizeBinaryProducer(BaseFixedWidthVector vector) { + super(vector); + reuseBytes = new byte[vector.getTypeWidth()]; + } + + @Override + public void produce(Encoder encoder) throws IOException { + long offset = (long) currentIndex * vector.getTypeWidth(); + vector.getDataBuffer().getBytes(offset, reuseBytes); + encoder.writeFixed(reuseBytes); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeListProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeListProducer.java new file mode 100644 index 0000000000..acb6fb8c00 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFixedSizeListProducer.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.FieldVector; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces array values from a {@link FixedSizeListVector}, writes data to an avro + * encoder. + */ +public class AvroFixedSizeListProducer extends BaseAvroProducer { + + private final Producer delegate; + + /** Instantiate an AvroFixedSizeListProducer. */ + public AvroFixedSizeListProducer( + FixedSizeListVector vector, Producer delegate) { + super(vector); + this.delegate = delegate; + } + + @Override + public void produce(Encoder encoder) throws IOException { + + encoder.writeArrayStart(); + encoder.setItemCount(vector.getListSize()); + + for (int i = 0; i < vector.getListSize(); i++) { + encoder.startItem(); + delegate.produce(encoder); + } + + encoder.writeArrayEnd(); + currentIndex++; + } + + @Override + public void skipNull() { + super.skipNull(); + // Child vector contains a fixed number of elements for each entry + int childIndex = currentIndex * vector.getListSize(); + delegate.setPosition(childIndex); + } + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + super.setPosition(index); + // Child vector contains a fixed number of elements for each entry + int childIndex = currentIndex * vector.getListSize(); + delegate.setPosition(childIndex); + } + + @Override + @SuppressWarnings("unchecked") + public void resetValueVector(FixedSizeListVector vector) { + ((Producer) delegate).resetValueVector(vector.getDataVector()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat2Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat2Producer.java new file mode 100644 index 0000000000..07e5ea3591 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat2Producer.java @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.memory.util.Float16; +import org.apache.arrow.vector.Float2Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces float values from a {@link Float2Vector}, writes data to an Avro encoder. + */ +public class AvroFloat2Producer extends BaseAvroProducer { + + /** Instantiate an AvroFloat2Producer. */ + public AvroFloat2Producer(Float2Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + short rawValue = vector.getDataBuffer().getShort(currentIndex * (long) Float2Vector.TYPE_WIDTH); + encoder.writeFloat(Float16.toFloat(rawValue)); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat4Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat4Producer.java new file mode 100644 index 0000000000..5121ba3a11 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat4Producer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.Float4Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces float values from a {@link Float4Vector}, writes data to an Avro encoder. + */ +public class AvroFloat4Producer extends BaseAvroProducer { + + /** Instantiate an AvroFloat4Producer. */ + public AvroFloat4Producer(Float4Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + float value = vector.getDataBuffer().getFloat(currentIndex * (long) Float4Vector.TYPE_WIDTH); + encoder.writeFloat(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat8Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat8Producer.java new file mode 100644 index 0000000000..05fca750ba --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroFloat8Producer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.Float8Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces double values from a {@link Float8Vector}, writes data to an Avro encoder. + */ +public class AvroFloat8Producer extends BaseAvroProducer { + + /** Instantiate an AvroFloat8Producer. */ + public AvroFloat8Producer(Float8Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + double value = vector.getDataBuffer().getDouble(currentIndex * (long) Float8Vector.TYPE_WIDTH); + encoder.writeDouble(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroIntProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroIntProducer.java new file mode 100644 index 0000000000..4c9cc9b71a --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroIntProducer.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.BaseFixedWidthVector; +import org.apache.arrow.vector.IntVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces int values from an {@link IntVector}, writes data to an avro encoder. + * + *

Logical types are also supported, for vectors derived from {@link BaseFixedWidthVector} where + * the internal representation matches IntVector and requires no conversion. + */ +public class AvroIntProducer extends BaseAvroProducer { + + /** Instantiate an AvroIntConsumer. */ + public AvroIntProducer(IntVector vector) { + super(vector); + } + + /** Protected constructor for a logical types with an integer representation. */ + protected AvroIntProducer(BaseFixedWidthVector vector) { + super(vector); + if (vector.getTypeWidth() != IntVector.TYPE_WIDTH) { + throw new IllegalArgumentException( + "AvroIntProducer requires type width = " + IntVector.TYPE_WIDTH); + } + } + + @Override + public void produce(Encoder encoder) throws IOException { + int value = vector.getDataBuffer().getInt(currentIndex * (long) IntVector.TYPE_WIDTH); + encoder.writeInt(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroListProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroListProducer.java new file mode 100644 index 0000000000..10cfe9549a --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroListProducer.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.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces array values from a {@link ListVector}, writes data to an avro encoder. + */ +public class AvroListProducer extends BaseAvroProducer { + + private final Producer delegate; + + /** Instantiate an AvroListProducer. */ + public AvroListProducer(ListVector vector, Producer delegate) { + super(vector); + this.delegate = delegate; + } + + @Override + public void produce(Encoder encoder) throws IOException { + + int startOffset = vector.getOffsetBuffer().getInt(currentIndex * (long) Integer.BYTES); + int endOffset = vector.getOffsetBuffer().getInt((currentIndex + 1) * (long) Integer.BYTES); + int nItems = endOffset - startOffset; + + encoder.writeArrayStart(); + encoder.setItemCount(nItems); + + for (int i = 0; i < nItems; i++) { + encoder.startItem(); + delegate.produce(encoder); + } + + encoder.writeArrayEnd(); + currentIndex++; + } + + // Do not override skipNull(), delegate will not have an entry if the list is null + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + int delegateOffset = vector.getOffsetBuffer().getInt(index * (long) Integer.BYTES); + delegate.setPosition(delegateOffset); + super.setPosition(index); + } + + @Override + @SuppressWarnings("unchecked") + public void resetValueVector(ListVector vector) { + ((Producer) delegate).resetValueVector(vector.getDataVector()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroMapProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroMapProducer.java new file mode 100644 index 0000000000..568d5b62e4 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroMapProducer.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.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.avro.io.Encoder; + +/** Producer which produces map type values to avro encoder. Write the data to {@link MapVector}. */ +public class AvroMapProducer extends BaseAvroProducer { + + private final Producer delegate; + + /** Instantiate a AvroMapProducer. */ + public AvroMapProducer(MapVector vector, Producer delegate) { + super(vector); + this.delegate = delegate; + } + + @Override + public void produce(Encoder encoder) throws IOException { + + int startOffset = vector.getOffsetBuffer().getInt(currentIndex * (long) Integer.BYTES); + int endOffset = vector.getOffsetBuffer().getInt((currentIndex + 1) * (long) Integer.BYTES); + int nEntries = endOffset - startOffset; + + encoder.writeMapStart(); + encoder.setItemCount(nEntries); + + for (int i = 0; i < nEntries; i++) { + encoder.startItem(); + delegate.produce(encoder); + } + + encoder.writeMapEnd(); + currentIndex++; + } + + // Do not override skipNull(), delegate will not have an entry if the map is null + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + int delegateOffset = vector.getOffsetBuffer().getInt(index * (long) Integer.BYTES); + delegate.setPosition(delegateOffset); + super.setPosition(index); + } + + @Override + @SuppressWarnings("unchecked") + public void resetValueVector(MapVector vector) { + ((Producer) delegate).resetValueVector(vector.getDataVector()); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullProducer.java new file mode 100644 index 0000000000..1bd1e891f1 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullProducer.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.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.NullVector; +import org.apache.avro.io.Encoder; + +/** Producer that produces null values from a {@link NullVector}, writes data to an Avro encoder. */ +public class AvroNullProducer extends BaseAvroProducer { + + /** Instantiate an AvroNullProducer. */ + public AvroNullProducer(NullVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + encoder.writeNull(); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullableProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullableProducer.java new file mode 100644 index 0000000000..5f8b314f49 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroNullableProducer.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.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; +import org.apache.avro.io.Encoder; + +/** + * Producer wrapper which produces nullable types to an avro encoder. Read data from the underlying + * {@link FieldVector}. + * + * @param The vector within producer or its delegate, used for partially produce purpose. + */ +public class AvroNullableProducer extends BaseAvroProducer { + + private final Producer delegate; + + /** Instantiate a AvroNullableProducer. */ + public AvroNullableProducer(Producer delegate) { + super(delegate.getVector()); + this.delegate = delegate; + } + + @Override + public void produce(Encoder encoder) throws IOException { + if (vector.isNull(currentIndex)) { + encoder.writeInt(1); + encoder.writeNull(); + delegate.skipNull(); + } else { + encoder.writeInt(0); + delegate.produce(encoder); + } + currentIndex++; + } + + @Override + public void skipNull() { + // Can be called by containers of nullable types + delegate.skipNull(); + currentIndex++; + } + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + delegate.setPosition(index); + super.setPosition(index); + } + + @Override + public void resetValueVector(T vector) { + delegate.resetValueVector(vector); + } + + @Override + public T getVector() { + return delegate.getVector(); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroSmallIntProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroSmallIntProducer.java new file mode 100644 index 0000000000..9c37750d9f --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroSmallIntProducer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.SmallIntVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces int values from an {@link SmallIntVector}, writes data to an avro encoder. + */ +public class AvroSmallIntProducer extends BaseAvroProducer { + + /** Instantiate an AvroSmallIntProducer. */ + public AvroSmallIntProducer(SmallIntVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + short value = vector.getDataBuffer().getShort(currentIndex * (long) SmallIntVector.TYPE_WIDTH); + encoder.writeInt(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStringProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStringProducer.java new file mode 100644 index 0000000000..19e165cd13 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStringProducer.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.adapter.avro.producers; + +import java.io.IOException; +import java.nio.ByteBuffer; +import org.apache.arrow.vector.VarCharVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces string values from a {@link VarCharVector}, writes data to an Avro + * encoder. + */ +public class AvroStringProducer extends BaseAvroProducer { + + /** Instantiate an AvroStringProducer. */ + public AvroStringProducer(VarCharVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + + int start = vector.getStartOffset(currentIndex); + int end = vector.getEndOffset(currentIndex); + int length = end - start; + + // The nio ByteBuffer is created once per call, but underlying data is not copied + ByteBuffer nioBuffer = vector.getDataBuffer().nioBuffer(start, length); + encoder.writeBytes(nioBuffer); + + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStructProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStructProducer.java new file mode 100644 index 0000000000..86c1949bf6 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroStructProducer.java @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.FieldVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.avro.io.Encoder; + +/** + * Producer which produces nested record type values to avro encoder. Read the data from {@link + * org.apache.arrow.vector.complex.StructVector}. + */ +public class AvroStructProducer extends BaseAvroProducer { + + private final Producer[] delegates; + + /** Instantiate a AvroStructProducer. */ + public AvroStructProducer(StructVector vector, Producer[] delegates) { + super(vector); + this.delegates = delegates; + } + + @Override + public void produce(Encoder encoder) throws IOException { + + for (Producer delegate : delegates) { + delegate.produce(encoder); + } + + currentIndex++; + } + + @Override + public void skipNull() { + for (Producer delegate : delegates) { + delegate.skipNull(); + } + super.skipNull(); + } + + @Override + public void setPosition(int index) { + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds: " + index); + } + for (Producer delegate : delegates) { + delegate.setPosition(index); + } + super.setPosition(index); + } + + @Override + @SuppressWarnings("unchecked") + public void resetValueVector(StructVector vector) { + for (int i = 0; i < delegates.length; i++) { + Producer delegate = (Producer) delegates[i]; + delegate.resetValueVector(vector.getChildrenFromFields().get(i)); + } + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroTinyIntProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroTinyIntProducer.java new file mode 100644 index 0000000000..30a80e5094 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroTinyIntProducer.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.TinyIntVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces int values from an {@link TinyIntVector}, writes data to an avro encoder. + */ +public class AvroTinyIntProducer extends BaseAvroProducer { + + /** Instantiate an AvroTinyIntProducer. */ + public AvroTinyIntProducer(TinyIntVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + byte value = vector.getDataBuffer().getByte(currentIndex * (long) TinyIntVector.TYPE_WIDTH); + encoder.writeInt(value); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint1Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint1Producer.java new file mode 100644 index 0000000000..83cbc9ef8e --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint1Producer.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.UInt1Vector; +import org.apache.avro.io.Encoder; + +/** Producer that produces int values from a {@link UInt1Vector}, writes data to an avro encoder. */ +public class AvroUint1Producer extends BaseAvroProducer { + + /** Instantiate an AvroUint1Producer. */ + public AvroUint1Producer(UInt1Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + byte unsigned = vector.getDataBuffer().getByte(currentIndex * (long) UInt1Vector.TYPE_WIDTH); + int unsignedInt = Byte.toUnsignedInt(unsigned); + encoder.writeInt(unsignedInt); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint2Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint2Producer.java new file mode 100644 index 0000000000..1e30c82cd2 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint2Producer.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.UInt2Vector; +import org.apache.avro.io.Encoder; + +/** Producer that produces int values from a {@link UInt2Vector}, writes data to an avro encoder. */ +public class AvroUint2Producer extends BaseAvroProducer { + + /** Instantiate an AvroUint2Producer. */ + public AvroUint2Producer(UInt2Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + short unsigned = vector.getDataBuffer().getShort(currentIndex * (long) UInt2Vector.TYPE_WIDTH); + int unsignedInt = Short.toUnsignedInt(unsigned); + encoder.writeInt(unsignedInt); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint4Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint4Producer.java new file mode 100644 index 0000000000..63f78429dd --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint4Producer.java @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.UInt4Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces long values from a {@link UInt4Vector}, writes data to an avro encoder. + */ +public class AvroUint4Producer extends BaseAvroProducer { + + /** Instantiate an AvroUint4Producer. */ + public AvroUint4Producer(UInt4Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + int unsigned = vector.getDataBuffer().getInt(currentIndex * (long) UInt4Vector.TYPE_WIDTH); + long unsignedLong = Integer.toUnsignedLong(unsigned); + encoder.writeLong(unsignedLong); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint8Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint8Producer.java new file mode 100644 index 0000000000..819b4d4140 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/AvroUint8Producer.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.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.UInt8Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces long values from a {@link UInt8Vector}, writes data to an avro encoder. + */ +public class AvroUint8Producer extends BaseAvroProducer { + + /** Instantiate an AvroUint8Producer. */ + public AvroUint8Producer(UInt8Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + long unsigned = vector.getDataBuffer().getLong(currentIndex * (long) UInt8Vector.TYPE_WIDTH); + if (unsigned < 0) { + throw new ArithmeticException("Unsigned long value is too large for Avro encoding"); + } + encoder.writeLong(unsigned); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/BaseAvroProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/BaseAvroProducer.java new file mode 100644 index 0000000000..30c004bdc6 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/BaseAvroProducer.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.adapter.avro.producers; + +import org.apache.arrow.vector.FieldVector; + +/** + * Base class for avro producers. + * + * @param vector type. + */ +public abstract class BaseAvroProducer implements Producer { + + protected T vector; + protected int currentIndex; + + /** + * Constructs a base avro consumer. + * + * @param vector the vector to consume. + */ + protected BaseAvroProducer(T vector) { + this.vector = vector; + } + + @Override + public void skipNull() { + currentIndex++; + } + + /** + * Sets the current index for this producer against the underlying vector. + * + *

For a vector of length N, the valid range is [0, N] inclusive. Setting index = N signifies + * that no further data is available for production (this is the state the produce will be in when + * production for the current vector is complete). + * + * @param index New current index for the producer + */ + @Override + public void setPosition(int index) { + // currentIndex == value count is a valid state, no more values will be produced + if (index < 0 || index > vector.getValueCount()) { + throw new IllegalArgumentException("Index out of bounds"); + } + currentIndex = index; + } + + @Override + public void resetValueVector(T vector) { + this.vector = vector; + this.currentIndex = 0; + } + + @Override + public T getVector() { + return vector; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/CompositeAvroProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/CompositeAvroProducer.java new file mode 100644 index 0000000000..d1ed506108 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/CompositeAvroProducer.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import java.util.List; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.avro.io.Encoder; + +/** Composite producer which holds all producers. It manages the produce and cleanup process. */ +public class CompositeAvroProducer { + + private final List> producers; + + public CompositeAvroProducer(List> producers) { + this.producers = producers; + } + + public List> getProducers() { + return producers; + } + + /** Produce encoder data. */ + public void produce(Encoder encoder) throws IOException { + for (Producer producer : producers) { + producer.produce(encoder); + } + } + + /** Reset vector of consumers with the given {@link VectorSchemaRoot}. */ + @SuppressWarnings({"unchecked", "rawtypes"}) + public void resetProducerVectors(VectorSchemaRoot root) { + // This method assumes that the VSR matches the constructed set of producers + int index = 0; + for (Producer producer : producers) { + producer.resetValueVector(root.getFieldVectors().get(index)); + } + } +} 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/main/java/org/apache/arrow/adapter/avro/producers/Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/Producer.java new file mode 100644 index 0000000000..aed2543348 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/Producer.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers; + +import java.io.IOException; +import org.apache.arrow.vector.FieldVector; +import org.apache.avro.io.Encoder; + +/** + * Interface that is used to produce values to avro encoder. + * + * @param The vector within producer or its delegate, used for partially produce purpose. + */ +public interface Producer { + + /** + * Produce a specific type value from the vector and write it to avro encoder. + * + * @param encoder avro encoder to write data + * @throws IOException on error + */ + void produce(Encoder encoder) throws IOException; + + /** Skip null value in the vector by setting reader position + 1. */ + void skipNull(); + + /** Set the position to read value from vector. */ + void setPosition(int index); + + /** Reset the vector within producer. */ + void resetValueVector(T vector); + + /** Get the vector within the producer. */ + T getVector(); +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateDayProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateDayProducer.java new file mode 100644 index 0000000000..36680fb196 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateDayProducer.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.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroIntProducer; +import org.apache.arrow.vector.DateDayVector; + +/** + * Producer that produces date values from a {@link DateDayVector}, writes data to an Avro encoder. + */ +public class AvroDateDayProducer extends AvroIntProducer { + + // Date stored as integer number of days, matches Avro date type + + /** Instantiate an AvroDateProducer. */ + public AvroDateDayProducer(DateDayVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateMilliProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateMilliProducer.java new file mode 100644 index 0000000000..a64bb3a021 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDateMilliProducer.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.DateMilliVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that converts days in milliseconds from a {@link DateMilliVector} and produces date + * (INT) values, writes data to an Avro encoder. + */ +public class AvroDateMilliProducer extends BaseAvroProducer { + + // Convert milliseconds to days for Avro date type + + private static final long MILLIS_PER_DAY = 86400000; + + /** Instantiate an AvroDateMilliProducer. */ + public AvroDateMilliProducer(DateMilliVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + long millis = vector.getDataBuffer().getLong(currentIndex * (long) DateMilliVector.TYPE_WIDTH); + long days = millis / MILLIS_PER_DAY; + if (days > (long) Integer.MAX_VALUE || days < (long) Integer.MIN_VALUE) { + throw new ArithmeticException("Date value is too large for Avro encoding"); + } + encoder.writeInt((int) days); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimal256Producer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimal256Producer.java new file mode 100644 index 0000000000..f72aa6d9e0 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimal256Producer.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import java.io.IOException; +import java.math.BigDecimal; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.Decimal256Vector; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces decimal values from a {@link Decimal256Vector}, writes data to an Avro + * encoder. + */ +public class AvroDecimal256Producer extends BaseAvroProducer { + + // Logic is the same as for DecimalVector (128 bit) + + byte[] encodedBytes = new byte[Decimal256Vector.TYPE_WIDTH]; + + /** Instantiate an AvroDecimalProducer. */ + public AvroDecimal256Producer(Decimal256Vector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + BigDecimal value = vector.getObject(currentIndex++); + AvroDecimalProducer.encodeDecimal(value, encodedBytes); + encoder.writeFixed(encodedBytes); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimalProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimalProducer.java new file mode 100644 index 0000000000..51ad7c7200 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimalProducer.java @@ -0,0 +1,59 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.logical; + +import java.io.IOException; +import java.math.BigDecimal; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.util.DecimalUtility; +import org.apache.avro.io.Encoder; + +/** + * Producer that produces decimal values from a {@link DecimalVector}, writes data to an Avro + * encoder. + */ +public class AvroDecimalProducer extends BaseAvroProducer { + + // Arrow stores decimals with native endianness, but Avro requires big endian + // Writing the Arrow representation as fixed bytes fails on little-end machines + // Instead, we replicate the big endian logic explicitly here + // See DecimalUtility.writeByteArrayToArrowBufHelper + + byte[] encodedBytes = new byte[DecimalVector.TYPE_WIDTH]; + + /** Instantiate an AvroDecimalProducer. */ + public AvroDecimalProducer(DecimalVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + // Use getObject() to go back to a BigDecimal then re-encode + BigDecimal value = vector.getObject(currentIndex++); + encodeDecimal(value, encodedBytes); + encoder.writeFixed(encodedBytes); + } + + static void encodeDecimal(BigDecimal value, byte[] encodedBytes) { + byte[] valueBytes = value.unscaledValue().toByteArray(); + byte[] padding = valueBytes[0] < 0 ? DecimalUtility.minus_one : DecimalUtility.zeroes; + System.arraycopy(padding, 0, encodedBytes, 0, encodedBytes.length - valueBytes.length); + System.arraycopy( + valueBytes, 0, encodedBytes, encodedBytes.length - valueBytes.length, valueBytes.length); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMicroProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMicroProducer.java new file mode 100644 index 0000000000..203d102034 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMicroProducer.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeMicroVector; + +/** + * Producer that produces time (microseconds) values from a {@link TimeMicroVector}, writes data to + * an Avro encoder. + */ +public class AvroTimeMicroProducer extends AvroBigIntProducer { + + // Time in microseconds stored as long, matches Avro time-micros type + + /** Instantiate an AvroTimeMicroProducer. */ + public AvroTimeMicroProducer(TimeMicroVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMilliProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMilliProducer.java new file mode 100644 index 0000000000..2a452e75a4 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeMilliProducer.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroIntProducer; +import org.apache.arrow.vector.TimeMilliVector; + +/** + * Producer that produces time (milliseconds) values from a {@link TimeMilliVector}, writes data to + * an Avro encoder. + */ +public class AvroTimeMilliProducer extends AvroIntProducer { + + // Time in milliseconds stored as integer, matches Avro time-millis type + + /** Instantiate an AvroTimeMilliProducer. */ + public AvroTimeMilliProducer(TimeMilliVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeNanoProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeNanoProducer.java new file mode 100644 index 0000000000..7034dbbb50 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeNanoProducer.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.TimeNanoVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that converts nanoseconds from a {@link TimeNanoVector} and produces time (microseconds) + * values, writes data to an Avro encoder. + */ +public class AvroTimeNanoProducer extends BaseAvroProducer { + + // Convert nanoseconds to microseconds for Avro time-micros (LONG) type + // Range is 1000 times less than for microseconds, so the type will fit (with loss of precision) + + private static final long NANOS_PER_MICRO = 1000; + + public AvroTimeNanoProducer(TimeNanoVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + long nanos = vector.getDataBuffer().getLong(currentIndex * (long) TimeNanoVector.TYPE_WIDTH); + long micros = nanos / NANOS_PER_MICRO; + encoder.writeLong(micros); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeSecProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeSecProducer.java new file mode 100644 index 0000000000..951605b6c3 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimeSecProducer.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.TimeSecVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that converts seconds from a {@link TimeSecVector} and produces time (microseconds) + * values, writes data to an Avro encoder. + */ +public class AvroTimeSecProducer extends BaseAvroProducer { + + // Convert seconds to milliseconds for Avro time-millis (INT) type + // INT is enough to cover the number of milliseconds in a day + // So overflows should not happen if values are valid times of day + + private static final int MILLIS_PER_SECOND = 1000; + private static final long OVERFLOW_LIMIT = Integer.MAX_VALUE / 1000; + + /** Instantiate an AvroTimeSecProducer. */ + public AvroTimeSecProducer(TimeSecVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + int seconds = vector.getDataBuffer().getInt(currentIndex * (long) TimeSecVector.TYPE_WIDTH); + if (Math.abs(seconds) > OVERFLOW_LIMIT) { + throw new ArithmeticException("Time value is too large for Avro encoding"); + } + int millis = seconds * MILLIS_PER_SECOND; + encoder.writeInt(millis); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroProducer.java new file mode 100644 index 0000000000..4e744b5e76 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroProducer.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampMicroVector; + +/** + * Producer that produces local timestamp (microseconds) values from a {@link TimeStampMicroVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampMicroProducer extends AvroBigIntProducer { + + // Local timestamp in epoch microseconds stored as long, matches Avro local-timestamp-micros type + + /** Instantiate an AvroTimestampMicroProducer. */ + public AvroTimestampMicroProducer(TimeStampMicroVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroTzProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroTzProducer.java new file mode 100644 index 0000000000..ece7482303 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMicroTzProducer.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampMicroTZVector; + +/** + * Producer that produces UTC timestamp (microseconds) values from a {@link TimeStampMicroTZVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampMicroTzProducer extends AvroBigIntProducer { + + // UTC timestamp in epoch microseconds stored as long, matches Avro timestamp-micros type + // Both Arrow and Avro store zone-aware times in UTC so zone conversion is not needed + + /** Instantiate an AvroTimestampMicroTzProducer. */ + public AvroTimestampMicroTzProducer(TimeStampMicroTZVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliProducer.java new file mode 100644 index 0000000000..e71acff220 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliProducer.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampMilliVector; + +/** + * Producer that produces local timestamp (milliseconds) values from a {@link TimeStampMilliVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampMilliProducer extends AvroBigIntProducer { + + // Local timestamp in epoch milliseconds stored as long, matches Avro local-timestamp-millis type + + /** Instantiate an AvroTimestampMilliProducer. */ + public AvroTimestampMilliProducer(TimeStampMilliVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliTzProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliTzProducer.java new file mode 100644 index 0000000000..b1b55fca78 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampMilliTzProducer.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampMilliTZVector; + +/** + * Producer that produces UTC timestamp (milliseconds) values from a {@link TimeStampMilliTZVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampMilliTzProducer extends AvroBigIntProducer { + + // UTC timestamp in epoch milliseconds stored as long, matches Avro timestamp-millis type + // Both Arrow and Avro store zone-aware times in UTC so zone conversion is not needed + + /** Instantiate an AvroTimestampMilliTzProducer. */ + public AvroTimestampMilliTzProducer(TimeStampMilliTZVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoProducer.java new file mode 100644 index 0000000000..9e172ea91e --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoProducer.java @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampNanoVector; + +/** + * Producer that produces local timestamp (nanoseconds) values from a {@link TimeStampNanoVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampNanoProducer extends AvroBigIntProducer { + + // Local timestamp in epoch nanoseconds stored as long, matches Avro local-timestamp-nanos type + + /** Instantiate an AvroTimestampNanoProducer. */ + public AvroTimestampNanoProducer(TimeStampNanoVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoTzProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoTzProducer.java new file mode 100644 index 0000000000..ae261d8396 --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampNanoTzProducer.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.logical; + +import org.apache.arrow.adapter.avro.producers.AvroBigIntProducer; +import org.apache.arrow.vector.TimeStampNanoTZVector; + +/** + * Producer that produces local timestamp (nanoseconds) values from a {@link TimeStampNanoTZVector}, + * writes data to an Avro encoder. + */ +public class AvroTimestampNanoTzProducer extends AvroBigIntProducer { + + // UTC timestamp in epoch nanoseconds stored as long, matches Avro timestamp-nanos type + // Both Arrow and Avro store zone-aware times in UTC so zone conversion is not needed + + /** Instantiate an AvroTimestampNanoTzProducer. */ + public AvroTimestampNanoTzProducer(TimeStampNanoTZVector vector) { + super(vector); + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecProducer.java new file mode 100644 index 0000000000..a6ade2d19b --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecProducer.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adapter.avro.producers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that converts epoch seconds from a {@link TimeStampSecVector} and produces local + * timestamp (milliseconds) values, writes data to an Avro encoder. + */ +public class AvroTimestampSecProducer extends BaseAvroProducer { + + // Avro does not support timestamps in seconds, so convert to local-timestamp-millis type + // Check for overflow and raise an exception + + private static final long MILLIS_PER_SECOND = 1000; + private static final long OVERFLOW_LIMIT = Long.MAX_VALUE / MILLIS_PER_SECOND; + + /** Instantiate an AvroTimestampSecProducer. */ + public AvroTimestampSecProducer(TimeStampSecVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + long seconds = + vector.getDataBuffer().getLong(currentIndex * (long) TimeStampSecVector.TYPE_WIDTH); + if (Math.abs(seconds) > OVERFLOW_LIMIT) { + throw new ArithmeticException("Timestamp value is too large for Avro encoding"); + } + long millis = seconds * MILLIS_PER_SECOND; + encoder.writeLong(millis); + currentIndex++; + } +} diff --git a/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecTzProducer.java b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecTzProducer.java new file mode 100644 index 0000000000..bd6cc14dad --- /dev/null +++ b/adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroTimestampSecTzProducer.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.adapter.avro.producers.logical; + +import java.io.IOException; +import org.apache.arrow.adapter.avro.producers.BaseAvroProducer; +import org.apache.arrow.vector.TimeStampSecTZVector; +import org.apache.arrow.vector.TimeStampVector; +import org.apache.avro.io.Encoder; + +/** + * Producer that converts epoch seconds from a {@link TimeStampSecTZVector} and produces UTC + * timestamp (milliseconds) values, writes data to an Avro encoder. + */ +public class AvroTimestampSecTzProducer extends BaseAvroProducer { + + // Avro does not support timestamps in seconds, so convert to timestamp-millis type + // Check for overflow and raise an exception + + // Both Arrow and Avro store zone-aware times in UTC so zone conversion is not needed + + private static final long MILLIS_PER_SECOND = 1000; + private static final long OVERFLOW_LIMIT = Long.MAX_VALUE / MILLIS_PER_SECOND; + + /** Instantiate an AvroTimestampSecTzProducer. */ + public AvroTimestampSecTzProducer(TimeStampSecTZVector vector) { + super(vector); + } + + @Override + public void produce(Encoder encoder) throws IOException { + long utcSeconds = + vector.getDataBuffer().getLong(currentIndex * (long) TimeStampVector.TYPE_WIDTH); + if (Math.abs(utcSeconds) > OVERFLOW_LIMIT) { + throw new ArithmeticException("Timestamp value is too large for Avro encoding"); + } + long utcMillis = utcSeconds * MILLIS_PER_SECOND; + encoder.writeLong(utcMillis); + currentIndex++; + } +} 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 new file mode 100644 index 0000000000..6d66ee9d45 --- /dev/null +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java @@ -0,0 +1,2901 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +import static org.junit.jupiter.api.Assertions.*; + +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.io.InputStream; +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.nio.ByteBuffer; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZonedDateTime; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.util.Float16; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.Float2Vector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.NullVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeNanoVector; +import org.apache.arrow.vector.TimeSecVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; +import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.TimeStampNanoTZVector; +import org.apache.arrow.vector.TimeStampNanoVector; +import org.apache.arrow.vector.TimeStampSecTZVector; +import org.apache.arrow.vector.TimeStampSecVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.UInt1Vector; +import org.apache.arrow.vector.UInt2Vector; +import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.FixedSizeListVector; +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.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; +import org.apache.arrow.vector.util.JsonStringHashMap; +import org.apache.avro.Conversions; +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.EncoderFactory; +import org.apache.avro.util.Utf8; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class ArrowToAvroDataTest { + + @TempDir public static File TMP; + + // Data production for primitive types, nullable and non-nullable + + @Test + public void testWriteNullColumn() throws Exception { + + // Field definition + FieldType nullField = new FieldType(false, new ArrowType.Null(), null); + + // Create empty vector + NullVector nullVector = new NullVector(new Field("nullColumn", nullField, null)); + + int rowCount = 10; + + // Set up VSR + List vectors = Arrays.asList(nullVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set all values to null + for (int row = 0; row < rowCount; row++) { + nullVector.setNull(row); + } + + File dataFile = new File(TMP, "testWriteNullColumn.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); + 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()); + 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); + assertNull(record.get("nullColumn")); + } + } + } + } + + @Test + public void testWriteBooleans() throws Exception { + + // Field definition + FieldType booleanField = new FieldType(false, new ArrowType.Bool(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + BitVector booleanVector = new BitVector(new Field("boolean", booleanField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(booleanVector); + int rowCount = 10; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < rowCount; row++) { + booleanVector.set(row, row % 2 == 0 ? 1 : 0); + } + + File dataFile = new File(TMP, "testWriteBooleans.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); + 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()); + 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); + assertEquals(booleanVector.get(row) == 1, record.get("boolean")); + } + } + } + } + + @Test + public void testWriteNullableBooleans() throws Exception { + + // Field definition + FieldType booleanField = new FieldType(true, new ArrowType.Bool(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + BitVector booleanVector = new BitVector(new Field("boolean", booleanField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(booleanVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null value + booleanVector.setNull(0); + + // False value + booleanVector.set(1, 0); + + // True value + booleanVector.set(2, 1); + + File dataFile = new File(TMP, "testWriteNullableBooleans.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("boolean")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals(booleanVector.get(row) == 1, record.get("boolean")); + } + } + } + } + + @Test + public void testWriteIntegers() throws Exception { + + // Field definitions + FieldType int8Field = new FieldType(false, new ArrowType.Int(8, true), null); + FieldType int16Field = new FieldType(false, new ArrowType.Int(16, true), null); + FieldType int32Field = new FieldType(false, new ArrowType.Int(32, true), null); + FieldType int64Field = new FieldType(false, new ArrowType.Int(64, true), null); + FieldType uint8Field = new FieldType(false, new ArrowType.Int(8, false), null); + FieldType uint16Field = new FieldType(false, new ArrowType.Int(16, false), null); + FieldType uint32Field = new FieldType(false, new ArrowType.Int(32, false), null); + FieldType uint64Field = new FieldType(false, new ArrowType.Int(64, false), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TinyIntVector int8Vector = new TinyIntVector(new Field("int8", int8Field, null), allocator); + SmallIntVector int16Vector = + new SmallIntVector(new Field("int16", int16Field, null), allocator); + IntVector int32Vector = new IntVector(new Field("int32", int32Field, null), allocator); + BigIntVector int64Vector = new BigIntVector(new Field("int64", int64Field, null), allocator); + UInt1Vector uint8Vector = new UInt1Vector(new Field("uint8", uint8Field, null), allocator); + UInt2Vector uint16Vector = new UInt2Vector(new Field("uint16", uint16Field, null), allocator); + UInt4Vector uint32Vector = new UInt4Vector(new Field("uint32", uint32Field, null), allocator); + UInt8Vector uint64Vector = new UInt8Vector(new Field("uint64", uint64Field, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList( + int8Vector, + int16Vector, + int32Vector, + int64Vector, + uint8Vector, + uint16Vector, + uint32Vector, + uint64Vector); + + int rowCount = 12; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < 10; row++) { + int8Vector.set(row, 11 * row * (row % 2 == 0 ? 1 : -1)); + int16Vector.set(row, 63 * row * (row % 2 == 0 ? 1 : -1)); + int32Vector.set(row, 513 * row * (row % 2 == 0 ? 1 : -1)); + int64Vector.set(row, 3791L * row * (row % 2 == 0 ? 1 : -1)); + uint8Vector.set(row, 11 * row); + uint16Vector.set(row, 63 * row); + uint32Vector.set(row, 513 * row); + uint64Vector.set(row, 3791L * row); + } + + // Min values + int8Vector.set(10, Byte.MIN_VALUE); + int16Vector.set(10, Short.MIN_VALUE); + int32Vector.set(10, Integer.MIN_VALUE); + int64Vector.set(10, Long.MIN_VALUE); + uint8Vector.set(10, 0); + uint16Vector.set(10, 0); + uint32Vector.set(10, 0); + uint64Vector.set(10, 0); + + // Max values + int8Vector.set(11, Byte.MAX_VALUE); + int16Vector.set(11, Short.MAX_VALUE); + int32Vector.set(11, Integer.MAX_VALUE); + int64Vector.set(11, Long.MAX_VALUE); + uint8Vector.set(11, 0xff); + uint16Vector.set(11, 0xffff); + uint32Vector.set(11, 0xffffffff); + uint64Vector.set(11, Long.MAX_VALUE); // Max that can be encoded + + File dataFile = new File(TMP, "testWriteIntegers.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); + 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()); + 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); + assertEquals((int) int8Vector.get(row), record.get("int8")); + assertEquals((int) int16Vector.get(row), record.get("int16")); + assertEquals(int32Vector.get(row), record.get("int32")); + assertEquals(int64Vector.get(row), record.get("int64")); + assertEquals(Byte.toUnsignedInt(uint8Vector.get(row)), record.get("uint8")); + assertEquals(Short.toUnsignedInt((short) uint16Vector.get(row)), record.get("uint16")); + assertEquals(Integer.toUnsignedLong(uint32Vector.get(row)), record.get("uint32")); + assertEquals(uint64Vector.get(row), record.get("uint64")); + } + } + } + } + + @Test + public void testWriteNullableIntegers() throws Exception { + + // Field definitions + FieldType int8Field = new FieldType(true, new ArrowType.Int(8, true), null); + FieldType int16Field = new FieldType(true, new ArrowType.Int(16, true), null); + FieldType int32Field = new FieldType(true, new ArrowType.Int(32, true), null); + FieldType int64Field = new FieldType(true, new ArrowType.Int(64, true), null); + FieldType uint8Field = new FieldType(true, new ArrowType.Int(8, false), null); + FieldType uint16Field = new FieldType(true, new ArrowType.Int(16, false), null); + FieldType uint32Field = new FieldType(true, new ArrowType.Int(32, false), null); + FieldType uint64Field = new FieldType(true, new ArrowType.Int(64, false), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TinyIntVector int8Vector = new TinyIntVector(new Field("int8", int8Field, null), allocator); + SmallIntVector int16Vector = + new SmallIntVector(new Field("int16", int16Field, null), allocator); + IntVector int32Vector = new IntVector(new Field("int32", int32Field, null), allocator); + BigIntVector int64Vector = new BigIntVector(new Field("int64", int64Field, null), allocator); + UInt1Vector uint8Vector = new UInt1Vector(new Field("uint8", uint8Field, null), allocator); + UInt2Vector uint16Vector = new UInt2Vector(new Field("uint16", uint16Field, null), allocator); + UInt4Vector uint32Vector = new UInt4Vector(new Field("uint32", uint32Field, null), allocator); + UInt8Vector uint64Vector = new UInt8Vector(new Field("uint64", uint64Field, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList( + int8Vector, + int16Vector, + int32Vector, + int64Vector, + uint8Vector, + uint16Vector, + uint32Vector, + uint64Vector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null values + int8Vector.setNull(0); + int16Vector.setNull(0); + int32Vector.setNull(0); + int64Vector.setNull(0); + uint8Vector.setNull(0); + uint16Vector.setNull(0); + uint32Vector.setNull(0); + uint64Vector.setNull(0); + + // Zero values + int8Vector.set(1, 0); + int16Vector.set(1, 0); + int32Vector.set(1, 0); + int64Vector.set(1, 0); + uint8Vector.set(1, 0); + uint16Vector.set(1, 0); + uint32Vector.set(1, 0); + uint64Vector.set(1, 0); + + // Non-zero values + int8Vector.set(2, Byte.MAX_VALUE); + int16Vector.set(2, Short.MAX_VALUE); + int32Vector.set(2, Integer.MAX_VALUE); + int64Vector.set(2, Long.MAX_VALUE); + uint8Vector.set(2, Byte.MAX_VALUE); + uint16Vector.set(2, Short.MAX_VALUE); + uint32Vector.set(2, Integer.MAX_VALUE); + uint64Vector.set(2, Long.MAX_VALUE); + + File dataFile = new File(TMP, "testWriteNullableIntegers.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("int8")); + assertNull(record.get("int16")); + assertNull(record.get("int32")); + assertNull(record.get("int64")); + assertNull(record.get("uint8")); + assertNull(record.get("uint16")); + assertNull(record.get("uint32")); + assertNull(record.get("uint64")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals((int) int8Vector.get(row), record.get("int8")); + assertEquals((int) int16Vector.get(row), record.get("int16")); + assertEquals(int32Vector.get(row), record.get("int32")); + assertEquals(int64Vector.get(row), record.get("int64")); + assertEquals(Byte.toUnsignedInt(uint8Vector.get(row)), record.get("uint8")); + assertEquals(Short.toUnsignedInt((short) uint16Vector.get(row)), record.get("uint16")); + assertEquals(Integer.toUnsignedLong(uint32Vector.get(row)), record.get("uint32")); + assertEquals(uint64Vector.get(row), record.get("uint64")); + } + } + } + } + + @Test + public void testWriteFloatingPoints() throws Exception { + + // Field definitions + FieldType float16Field = + new FieldType(false, new ArrowType.FloatingPoint(FloatingPointPrecision.HALF), null); + FieldType float32Field = + new FieldType(false, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null); + FieldType float64Field = + new FieldType(false, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + Float2Vector float16Vector = + new Float2Vector(new Field("float16", float16Field, null), allocator); + Float4Vector float32Vector = + new Float4Vector(new Field("float32", float32Field, null), allocator); + Float8Vector float64Vector = + new Float8Vector(new Field("float64", float64Field, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(float16Vector, float32Vector, float64Vector); + int rowCount = 15; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < 10; row++) { + float16Vector.set(row, Float16.toFloat16(3.6f * row * (row % 2 == 0 ? 1.0f : -1.0f))); + float32Vector.set(row, 37.6f * row * (row % 2 == 0 ? 1 : -1)); + float64Vector.set(row, 37.6d * row * (row % 2 == 0 ? 1 : -1)); + } + + float16Vector.set(10, Float16.toFloat16(Float.MIN_VALUE)); + float32Vector.set(10, Float.MIN_VALUE); + float64Vector.set(10, Double.MIN_VALUE); + + float16Vector.set(11, Float16.toFloat16(Float.MAX_VALUE)); + float32Vector.set(11, Float.MAX_VALUE); + float64Vector.set(11, Double.MAX_VALUE); + + float16Vector.set(12, Float16.toFloat16(Float.NaN)); + float32Vector.set(12, Float.NaN); + float64Vector.set(12, Double.NaN); + + float16Vector.set(13, Float16.toFloat16(Float.POSITIVE_INFINITY)); + float32Vector.set(13, Float.POSITIVE_INFINITY); + float64Vector.set(13, Double.POSITIVE_INFINITY); + + float16Vector.set(14, Float16.toFloat16(Float.NEGATIVE_INFINITY)); + float32Vector.set(14, Float.NEGATIVE_INFINITY); + float64Vector.set(14, Double.NEGATIVE_INFINITY); + + File dataFile = new File(TMP, "testWriteFloatingPoints.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); + 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()); + 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); + assertEquals(float16Vector.getValueAsFloat(row), record.get("float16")); + assertEquals(float32Vector.get(row), record.get("float32")); + assertEquals(float64Vector.get(row), record.get("float64")); + } + } + } + } + + @Test + public void testWriteNullableFloatingPoints() throws Exception { + + // Field definitions + FieldType float16Field = + new FieldType(true, new ArrowType.FloatingPoint(FloatingPointPrecision.HALF), null); + FieldType float32Field = + new FieldType(true, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null); + FieldType float64Field = + new FieldType(true, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + Float2Vector float16Vector = + new Float2Vector(new Field("float16", float16Field, null), allocator); + Float4Vector float32Vector = + new Float4Vector(new Field("float32", float32Field, null), allocator); + Float8Vector float64Vector = + new Float8Vector(new Field("float64", float64Field, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(float16Vector, float32Vector, float64Vector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null values + float16Vector.setNull(0); + float32Vector.setNull(0); + float64Vector.setNull(0); + + // Zero values + float16Vector.setSafeWithPossibleTruncate(1, 0.0f); + float32Vector.set(1, 0.0f); + float64Vector.set(1, 0.0); + + // Non-zero values + float16Vector.setSafeWithPossibleTruncate(2, 1.0f); + float32Vector.set(2, 1.0f); + float64Vector.set(2, 1.0); + + File dataFile = new File(TMP, "testWriteNullableFloatingPoints.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("float16")); + assertNull(record.get("float32")); + assertNull(record.get("float64")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals(float16Vector.getValueAsFloat(row), record.get("float16")); + assertEquals(float32Vector.get(row), record.get("float32")); + assertEquals(float64Vector.get(row), record.get("float64")); + } + } + } + } + + @Test + public void testWriteStrings() throws Exception { + + // Field definition + FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarCharVector stringVector = + new VarCharVector(new Field("string", stringField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(stringVector); + int rowCount = 5; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + stringVector.setSafe(0, "Hello world!".getBytes()); + stringVector.setSafe(1, "<%**\r\n\t\\abc\0$$>".getBytes()); + stringVector.setSafe(2, "你好世界".getBytes()); + stringVector.setSafe(3, "مرحبا بالعالم".getBytes()); + stringVector.setSafe(4, "(P ∧ P ⇒ Q) ⇒ Q".getBytes()); + + File dataFile = new File(TMP, "testWriteStrings.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); + 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()); + 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); + assertEquals(stringVector.getObject(row).toString(), record.get("string").toString()); + } + } + } + } + + @Test + public void testWriteNullableStrings() throws Exception { + + // Field definition + FieldType stringField = new FieldType(true, new ArrowType.Utf8(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarCharVector stringVector = + new VarCharVector(new Field("string", stringField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(stringVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + stringVector.setNull(0); + stringVector.setSafe(1, "".getBytes()); + stringVector.setSafe(2, "not empty".getBytes()); + + File dataFile = new File(TMP, "testWriteNullableStrings.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("string")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals(stringVector.getObject(row).toString(), record.get("string").toString()); + } + } + } + } + + @Test + public void testWriteBinary() throws Exception { + + // Field definition + FieldType binaryField = new FieldType(false, new ArrowType.Binary(), null); + FieldType fixedField = new FieldType(false, new ArrowType.FixedSizeBinary(5), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarBinaryVector binaryVector = + new VarBinaryVector(new Field("binary", binaryField, null), allocator); + FixedSizeBinaryVector fixedVector = + new FixedSizeBinaryVector(new Field("fixed", fixedField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(binaryVector, fixedVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + binaryVector.setSafe(0, new byte[] {1, 2, 3}); + binaryVector.setSafe(1, new byte[] {4, 5, 6, 7}); + binaryVector.setSafe(2, new byte[] {8, 9}); + + fixedVector.setSafe(0, new byte[] {1, 2, 3, 4, 5}); + fixedVector.setSafe(1, new byte[] {4, 5, 6, 7, 8, 9}); + fixedVector.setSafe(2, new byte[] {8, 9, 10, 11, 12}); + + File dataFile = new File(TMP, "testWriteBinary.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); + 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()); + 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); + ByteBuffer buf = ((ByteBuffer) record.get("binary")); + byte[] bytes = new byte[buf.remaining()]; + buf.get(bytes); + byte[] fixedBytes = ((GenericData.Fixed) record.get("fixed")).bytes(); + assertArrayEquals(binaryVector.getObject(row), bytes); + assertArrayEquals(fixedVector.getObject(row), fixedBytes); + } + } + } + } + + @Test + public void testWriteNullableBinary() throws Exception { + + // Field definition + FieldType binaryField = new FieldType(true, new ArrowType.Binary(), null); + FieldType fixedField = new FieldType(true, new ArrowType.FixedSizeBinary(5), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarBinaryVector binaryVector = + new VarBinaryVector(new Field("binary", binaryField, null), allocator); + FixedSizeBinaryVector fixedVector = + new FixedSizeBinaryVector(new Field("fixed", fixedField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(binaryVector, fixedVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + binaryVector.setNull(0); + binaryVector.setSafe(1, new byte[] {}); + binaryVector.setSafe(2, new byte[] {10, 11, 12}); + + fixedVector.setNull(0); + fixedVector.setSafe(1, new byte[] {0, 0, 0, 0, 0}); + fixedVector.setSafe(2, new byte[] {10, 11, 12, 13, 14}); + + File dataFile = new File(TMP, "testWriteNullableBinary.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("binary")); + assertNull(record.get("fixed")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + ByteBuffer buf = ((ByteBuffer) record.get("binary")); + byte[] bytes = new byte[buf.remaining()]; + buf.get(bytes); + byte[] fixedBytes = ((GenericData.Fixed) record.get("fixed")).bytes(); + assertArrayEquals(binaryVector.getObject(row), bytes); + assertArrayEquals(fixedVector.getObject(row), fixedBytes); + } + } + } + } + + // Data production for logical types, nullable and non-nullable + + @Test + public void testWriteDecimals() throws Exception { + + // Field definitions + FieldType decimal128Field1 = new FieldType(false, new ArrowType.Decimal(38, 10, 128), null); + FieldType decimal128Field2 = new FieldType(false, new ArrowType.Decimal(38, 5, 128), null); + FieldType decimal256Field1 = new FieldType(false, new ArrowType.Decimal(76, 20, 256), null); + FieldType decimal256Field2 = new FieldType(false, new ArrowType.Decimal(76, 10, 256), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DecimalVector decimal128Vector1 = + new DecimalVector(new Field("decimal128_1", decimal128Field1, null), allocator); + DecimalVector decimal128Vector2 = + new DecimalVector(new Field("decimal128_2", decimal128Field2, null), allocator); + Decimal256Vector decimal256Vector1 = + new Decimal256Vector(new Field("decimal256_1", decimal256Field1, null), allocator); + Decimal256Vector decimal256Vector2 = + new Decimal256Vector(new Field("decimal256_2", decimal256Field2, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList(decimal128Vector1, decimal128Vector2, decimal256Vector1, decimal256Vector2); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + decimal128Vector1.setSafe( + 0, new BigDecimal("12345.67890").setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 1, new BigDecimal("-98765.43210").setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 2, new BigDecimal("54321.09876").setScale(10, RoundingMode.UNNECESSARY)); + + decimal128Vector2.setSafe( + 0, new BigDecimal("12345.67890").setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 1, new BigDecimal("-98765.43210").setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 2, new BigDecimal("54321.09876").setScale(5, RoundingMode.UNNECESSARY)); + + decimal256Vector1.setSafe( + 0, + new BigDecimal("12345678901234567890.12345678901234567890") + .setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 1, + new BigDecimal("-98765432109876543210.98765432109876543210") + .setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 2, + new BigDecimal("54321098765432109876.54321098765432109876") + .setScale(20, RoundingMode.UNNECESSARY)); + + decimal256Vector2.setSafe( + 0, + new BigDecimal("12345678901234567890.1234567890").setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 1, + new BigDecimal("-98765432109876543210.9876543210") + .setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 2, + new BigDecimal("54321098765432109876.5432109876").setScale(10, RoundingMode.UNNECESSARY)); + + File dataFile = new File(TMP, "testWriteDecimals.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); + 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()); + 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); + assertEquals( + decimal128Vector1.getObject(row), decodeFixedDecimal(record, "decimal128_1")); + assertEquals( + decimal128Vector2.getObject(row), decodeFixedDecimal(record, "decimal128_2")); + assertEquals( + decimal256Vector1.getObject(row), decodeFixedDecimal(record, "decimal256_1")); + assertEquals( + decimal256Vector2.getObject(row), decodeFixedDecimal(record, "decimal256_2")); + } + } + } + } + + @Test + public void testWriteNullableDecimals() throws Exception { + + // Field definitions + FieldType decimal128Field1 = new FieldType(true, new ArrowType.Decimal(38, 10, 128), null); + FieldType decimal128Field2 = new FieldType(true, new ArrowType.Decimal(38, 5, 128), null); + FieldType decimal256Field1 = new FieldType(true, new ArrowType.Decimal(76, 20, 256), null); + FieldType decimal256Field2 = new FieldType(true, new ArrowType.Decimal(76, 10, 256), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DecimalVector decimal128Vector1 = + new DecimalVector(new Field("decimal128_1", decimal128Field1, null), allocator); + DecimalVector decimal128Vector2 = + new DecimalVector(new Field("decimal128_2", decimal128Field2, null), allocator); + Decimal256Vector decimal256Vector1 = + new Decimal256Vector(new Field("decimal256_1", decimal256Field1, null), allocator); + Decimal256Vector decimal256Vector2 = + new Decimal256Vector(new Field("decimal256_2", decimal256Field2, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList(decimal128Vector1, decimal128Vector2, decimal256Vector1, decimal256Vector2); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + decimal128Vector1.setNull(0); + decimal128Vector1.setSafe(1, BigDecimal.ZERO.setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 2, new BigDecimal("12345.67890").setScale(10, RoundingMode.UNNECESSARY)); + + decimal128Vector2.setNull(0); + decimal128Vector2.setSafe(1, BigDecimal.ZERO.setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 2, new BigDecimal("98765.43210").setScale(5, RoundingMode.UNNECESSARY)); + + decimal256Vector1.setNull(0); + decimal256Vector1.setSafe(1, BigDecimal.ZERO.setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 2, + new BigDecimal("12345678901234567890.12345678901234567890") + .setScale(20, RoundingMode.UNNECESSARY)); + + decimal256Vector2.setNull(0); + decimal256Vector2.setSafe(1, BigDecimal.ZERO.setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 2, + new BigDecimal("98765432109876543210.9876543210").setScale(10, RoundingMode.UNNECESSARY)); + + File dataFile = new File(TMP, "testWriteNullableDecimals.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("decimal128_1")); + assertNull(record.get("decimal128_2")); + assertNull(record.get("decimal256_1")); + assertNull(record.get("decimal256_2")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals( + decimal128Vector1.getObject(row), decodeFixedDecimal(record, "decimal128_1")); + assertEquals( + decimal128Vector2.getObject(row), decodeFixedDecimal(record, "decimal128_2")); + assertEquals( + decimal256Vector1.getObject(row), decodeFixedDecimal(record, "decimal256_1")); + assertEquals( + decimal256Vector2.getObject(row), decodeFixedDecimal(record, "decimal256_2")); + } + } + } + } + + private static BigDecimal decodeFixedDecimal(GenericRecord record, String fieldName) { + GenericData.Fixed fixed = (GenericData.Fixed) record.get(fieldName); + var logicalType = LogicalTypes.fromSchema(fixed.getSchema()); + return new Conversions.DecimalConversion().fromFixed(fixed, fixed.getSchema(), logicalType); + } + + @Test + public void testWriteDates() throws Exception { + + // Field definitions + FieldType dateDayField = new FieldType(false, new ArrowType.Date(DateUnit.DAY), null); + FieldType dateMillisField = + new FieldType(false, new ArrowType.Date(DateUnit.MILLISECOND), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DateDayVector dateDayVector = + new DateDayVector(new Field("dateDay", dateDayField, null), allocator); + DateMilliVector dateMillisVector = + new DateMilliVector(new Field("dateMillis", dateMillisField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(dateDayVector, dateMillisVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + dateDayVector.setSafe(0, (int) LocalDate.now().toEpochDay()); + dateDayVector.setSafe(1, (int) LocalDate.now().toEpochDay() + 1); + dateDayVector.setSafe(2, (int) LocalDate.now().toEpochDay() + 2); + + dateMillisVector.setSafe(0, LocalDate.now().toEpochDay() * 86400000L); + dateMillisVector.setSafe(1, (LocalDate.now().toEpochDay() + 1) * 86400000L); + dateMillisVector.setSafe(2, (LocalDate.now().toEpochDay() + 2) * 86400000L); + + File dataFile = new File(TMP, "testWriteDates.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); + 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()); + 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); + assertEquals(dateDayVector.get(row), record.get("dateDay")); + assertEquals( + dateMillisVector.get(row), ((long) (Integer) record.get("dateMillis")) * 86400000L); + } + } + } + } + + @Test + public void testWriteNullableDates() throws Exception { + + // Field definitions + FieldType dateDayField = new FieldType(true, new ArrowType.Date(DateUnit.DAY), null); + FieldType dateMillisField = new FieldType(true, new ArrowType.Date(DateUnit.MILLISECOND), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DateDayVector dateDayVector = + new DateDayVector(new Field("dateDay", dateDayField, null), allocator); + DateMilliVector dateMillisVector = + new DateMilliVector(new Field("dateMillis", dateMillisField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(dateDayVector, dateMillisVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + dateDayVector.setNull(0); + dateDayVector.setSafe(1, 0); + dateDayVector.setSafe(2, (int) LocalDate.now().toEpochDay()); + + dateMillisVector.setNull(0); + dateMillisVector.setSafe(1, 0); + dateMillisVector.setSafe(2, LocalDate.now().toEpochDay() * 86400000L); + + File dataFile = new File(TMP, "testWriteNullableDates.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("dateDay")); + assertNull(record.get("dateMillis")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals(dateDayVector.get(row), record.get("dateDay")); + assertEquals( + dateMillisVector.get(row), ((long) (Integer) record.get("dateMillis")) * 86400000L); + } + } + } + } + + @Test + public void testWriteTimes() throws Exception { + + // Field definitions + FieldType timeSecField = new FieldType(false, new ArrowType.Time(TimeUnit.SECOND, 32), null); + FieldType timeMillisField = + new FieldType(false, new ArrowType.Time(TimeUnit.MILLISECOND, 32), null); + FieldType timeMicrosField = + new FieldType(false, new ArrowType.Time(TimeUnit.MICROSECOND, 64), null); + FieldType timeNanosField = + new FieldType(false, new ArrowType.Time(TimeUnit.NANOSECOND, 64), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeSecVector timeSecVector = + new TimeSecVector(new Field("timeSec", timeSecField, null), allocator); + TimeMilliVector timeMillisVector = + new TimeMilliVector(new Field("timeMillis", timeMillisField, null), allocator); + TimeMicroVector timeMicrosVector = + new TimeMicroVector(new Field("timeMicros", timeMicrosField, null), allocator); + TimeNanoVector timeNanosVector = + new TimeNanoVector(new Field("timeNanos", timeNanosField, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList(timeSecVector, timeMillisVector, timeMicrosVector, timeNanosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timeSecVector.setSafe(0, ZonedDateTime.now().toLocalTime().toSecondOfDay()); + timeSecVector.setSafe(1, ZonedDateTime.now().toLocalTime().toSecondOfDay() - 1); + timeSecVector.setSafe(2, ZonedDateTime.now().toLocalTime().toSecondOfDay() - 2); + + timeMillisVector.setSafe( + 0, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000)); + timeMillisVector.setSafe( + 1, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000) - 1000); + timeMillisVector.setSafe( + 2, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000) - 2000); + + timeMicrosVector.setSafe(0, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000); + timeMicrosVector.setSafe(1, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000 - 1000000); + timeMicrosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000 - 2000000); + + timeNanosVector.setSafe(0, ZonedDateTime.now().toLocalTime().toNanoOfDay()); + timeNanosVector.setSafe(1, ZonedDateTime.now().toLocalTime().toNanoOfDay() - 1000000000); + timeNanosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay() - 2000000000); + + File dataFile = new File(TMP, "testWriteTimes.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); + 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()); + 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); + assertEquals(timeSecVector.get(row), (int) (record.get("timeSec")) / 1000); + assertEquals(timeMillisVector.get(row), record.get("timeMillis")); + assertEquals(timeMicrosVector.get(row), record.get("timeMicros")); + // Avro doesn't have time-nanos (mar 2025), so expect column to be saved as micros + long nanosAsMicros = (timeNanosVector.get(row) / 1000); + assertEquals(nanosAsMicros, (long) record.get("timeNanos")); + } + } + } + } + + @Test + public void testWriteNullableTimes() throws Exception { + + // Field definitions + FieldType timeSecField = new FieldType(true, new ArrowType.Time(TimeUnit.SECOND, 32), null); + FieldType timeMillisField = + new FieldType(true, new ArrowType.Time(TimeUnit.MILLISECOND, 32), null); + FieldType timeMicrosField = + new FieldType(true, new ArrowType.Time(TimeUnit.MICROSECOND, 64), null); + FieldType timeNanosField = + new FieldType(true, new ArrowType.Time(TimeUnit.NANOSECOND, 64), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeSecVector timeSecVector = + new TimeSecVector(new Field("timeSec", timeSecField, null), allocator); + TimeMilliVector timeMillisVector = + new TimeMilliVector(new Field("timeMillis", timeMillisField, null), allocator); + TimeMicroVector timeMicrosVector = + new TimeMicroVector(new Field("timeMicros", timeMicrosField, null), allocator); + TimeNanoVector timeNanosVector = + new TimeNanoVector(new Field("timeNanos", timeNanosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList(timeSecVector, timeMillisVector, timeMicrosVector, timeNanosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timeSecVector.setNull(0); + timeSecVector.setSafe(1, 0); + timeSecVector.setSafe(2, ZonedDateTime.now().toLocalTime().toSecondOfDay()); + + timeMillisVector.setNull(0); + timeMillisVector.setSafe(1, 0); + timeMillisVector.setSafe( + 2, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000)); + + timeMicrosVector.setNull(0); + timeMicrosVector.setSafe(1, 0); + timeMicrosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000); + + timeNanosVector.setNull(0); + timeNanosVector.setSafe(1, 0); + timeNanosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay()); + + File dataFile = new File(TMP, "testWriteNullableTimes.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("timeSec")); + assertNull(record.get("timeMillis")); + assertNull(record.get("timeMicros")); + assertNull(record.get("timeNanos")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals(timeSecVector.get(row), ((int) record.get("timeSec") / 1000)); + assertEquals(timeMillisVector.get(row), record.get("timeMillis")); + assertEquals(timeMicrosVector.get(row), record.get("timeMicros")); + // Avro doesn't have time-nanos (mar 2025), so expect column to be saved as micros + long nanosAsMicros = (timeNanosVector.get(row) / 1000); + assertEquals(nanosAsMicros, (long) record.get("timeNanos")); + } + } + } + } + + @Test + public void testWriteZoneAwareTimestamps() throws Exception { + + // Field definitions + FieldType timestampSecField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"), null); + FieldType timestampMillisField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"), null); + FieldType timestampMicrosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), null); + FieldType timestampNanosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampSecTZVector timestampSecVector = + new TimeStampSecTZVector(new Field("timestampSec", timestampSecField, null), allocator); + TimeStampMilliTZVector timestampMillisVector = + new TimeStampMilliTZVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroTZVector timestampMicrosVector = + new TimeStampMicroTZVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoTZVector timestampNanosVector = + new TimeStampNanoTZVector( + new Field("timestampNanos", timestampNanosField, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList( + timestampSecVector, timestampMillisVector, timestampMicrosVector, timestampNanosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampSecVector.setSafe(0, (int) Instant.now().getEpochSecond()); + timestampSecVector.setSafe(1, (int) Instant.now().getEpochSecond() - 1); + timestampSecVector.setSafe(2, (int) Instant.now().getEpochSecond() - 2); + + timestampMillisVector.setSafe(0, (int) Instant.now().toEpochMilli()); + timestampMillisVector.setSafe(1, (int) Instant.now().toEpochMilli() - 1000); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli() - 2000); + + timestampMicrosVector.setSafe(0, Instant.now().toEpochMilli() * 1000); + timestampMicrosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000); + timestampMicrosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000); + + timestampNanosVector.setSafe(0, Instant.now().toEpochMilli() * 1000000); + timestampNanosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000000); + timestampNanosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000000); + + File dataFile = new File(TMP, "testWriteZoneAwareTimestamps.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); + 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()); + 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); + assertEquals( + timestampSecVector.get(row), (int) ((long) record.get("timestampSec") / 1000)); + assertEquals(timestampMillisVector.get(row), (int) (long) record.get("timestampMillis")); + assertEquals(timestampMicrosVector.get(row), record.get("timestampMicros")); + assertEquals(timestampNanosVector.get(row), record.get("timestampNanos")); + } + } + } + } + + @Test + public void testWriteNullableZoneAwareTimestamps() throws Exception { + + // Field definitions + FieldType timestampSecField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"), null); + FieldType timestampMillisField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"), null); + FieldType timestampMicrosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), null); + FieldType timestampNanosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampSecTZVector timestampSecVector = + new TimeStampSecTZVector(new Field("timestampSec", timestampSecField, null), allocator); + TimeStampMilliTZVector timestampMillisVector = + new TimeStampMilliTZVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroTZVector timestampMicrosVector = + new TimeStampMicroTZVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoTZVector timestampNanosVector = + new TimeStampNanoTZVector( + new Field("timestampNanos", timestampNanosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList( + timestampSecVector, timestampMillisVector, timestampMicrosVector, timestampNanosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampSecVector.setNull(0); + timestampSecVector.setSafe(1, 0); + timestampSecVector.setSafe(2, (int) Instant.now().getEpochSecond()); + + timestampMillisVector.setNull(0); + timestampMillisVector.setSafe(1, 0); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli()); + + timestampMicrosVector.setNull(0); + timestampMicrosVector.setSafe(1, 0); + timestampMicrosVector.setSafe(2, Instant.now().toEpochMilli() * 1000); + + timestampNanosVector.setNull(0); + timestampNanosVector.setSafe(1, 0); + timestampNanosVector.setSafe(2, Instant.now().toEpochMilli() * 1000000); + + File dataFile = new File(TMP, "testWriteNullableZoneAwareTimestamps.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("timestampSec")); + assertNull(record.get("timestampMillis")); + assertNull(record.get("timestampMicros")); + assertNull(record.get("timestampNanos")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals( + timestampSecVector.get(row), (int) ((long) record.get("timestampSec") / 1000)); + assertEquals(timestampMillisVector.get(row), (int) (long) record.get("timestampMillis")); + assertEquals(timestampMicrosVector.get(row), record.get("timestampMicros")); + assertEquals(timestampNanosVector.get(row), record.get("timestampNanos")); + } + } + } + } + + @Test + public void testWriteLocalTimestamps() throws Exception { + + // Field definitions + FieldType timestampSecField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.SECOND, null), null); + FieldType timestampMillisField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null); + FieldType timestampMicrosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), null); + FieldType timestampNanosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampSecVector timestampSecVector = + new TimeStampSecVector(new Field("timestampSec", timestampSecField, null), allocator); + TimeStampMilliVector timestampMillisVector = + new TimeStampMilliVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroVector timestampMicrosVector = + new TimeStampMicroVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoVector timestampNanosVector = + new TimeStampNanoVector(new Field("timestampNanos", timestampNanosField, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList( + timestampSecVector, timestampMillisVector, timestampMicrosVector, timestampNanosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampSecVector.setSafe(0, (int) Instant.now().getEpochSecond()); + timestampSecVector.setSafe(1, (int) Instant.now().getEpochSecond() - 1); + timestampSecVector.setSafe(2, (int) Instant.now().getEpochSecond() - 2); + + timestampMillisVector.setSafe(0, (int) Instant.now().toEpochMilli()); + timestampMillisVector.setSafe(1, (int) Instant.now().toEpochMilli() - 1000); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli() - 2000); + + timestampMicrosVector.setSafe(0, Instant.now().toEpochMilli() * 1000); + timestampMicrosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000); + timestampMicrosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000); + + timestampNanosVector.setSafe(0, Instant.now().toEpochMilli() * 1000000); + timestampNanosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000000); + timestampNanosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000000); + + File dataFile = new File(TMP, "testWriteZoneAwareTimestamps.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); + 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()); + 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); + assertEquals( + timestampSecVector.get(row), (int) ((long) record.get("timestampSec") / 1000)); + assertEquals(timestampMillisVector.get(row), (int) (long) record.get("timestampMillis")); + assertEquals(timestampMicrosVector.get(row), record.get("timestampMicros")); + assertEquals(timestampNanosVector.get(row), record.get("timestampNanos")); + } + } + } + } + + @Test + public void testWriteNullableLocalTimestamps() throws Exception { + + // Field definitions + FieldType timestampSecField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.SECOND, null), null); + FieldType timestampMillisField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null); + FieldType timestampMicrosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), null); + FieldType timestampNanosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampSecVector timestampSecVector = + new TimeStampSecVector(new Field("timestampSec", timestampSecField, null), allocator); + TimeStampMilliVector timestampMillisVector = + new TimeStampMilliVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroVector timestampMicrosVector = + new TimeStampMicroVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoVector timestampNanosVector = + new TimeStampNanoVector(new Field("timestampNanos", timestampNanosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList( + timestampSecVector, timestampMillisVector, timestampMicrosVector, timestampNanosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampSecVector.setNull(0); + timestampSecVector.setSafe(1, 0); + timestampSecVector.setSafe(2, (int) Instant.now().getEpochSecond()); + + timestampMillisVector.setNull(0); + timestampMillisVector.setSafe(1, 0); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli()); + + timestampMicrosVector.setNull(0); + timestampMicrosVector.setSafe(1, 0); + timestampMicrosVector.setSafe(2, Instant.now().toEpochMilli() * 1000); + + timestampNanosVector.setNull(0); + timestampNanosVector.setSafe(1, 0); + timestampNanosVector.setSafe(2, Instant.now().toEpochMilli() * 1000000); + + File dataFile = new File(TMP, "testWriteNullableZoneAwareTimestamps.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); + 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()); + GenericDatumReader datumReader = new GenericDatumReader<>(schema); + + try (InputStream inputStream = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + + // Read and check values + GenericRecord record = datumReader.read(null, decoder); + assertNull(record.get("timestampSec")); + assertNull(record.get("timestampMillis")); + assertNull(record.get("timestampMicros")); + assertNull(record.get("timestampNanos")); + + for (int row = 1; row < rowCount; row++) { + record = datumReader.read(record, decoder); + assertEquals( + timestampSecVector.get(row), (int) ((long) record.get("timestampSec") / 1000)); + assertEquals(timestampMillisVector.get(row), (int) (long) record.get("timestampMillis")); + assertEquals(timestampMicrosVector.get(row), record.get("timestampMicros")); + assertEquals(timestampNanosVector.get(row), record.get("timestampNanos")); + } + } + } + } + + // Data production for containers of primitive and logical types, nullable and non-nullable + + @Test + public void testWriteLists() throws Exception { + + // Field definitions + FieldType intListField = new FieldType(false, new ArrowType.List(), null); + FieldType stringListField = new FieldType(false, new ArrowType.List(), null); + FieldType dateListField = new FieldType(false, new ArrowType.List(), null); + + Field intField = new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("item", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("item", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + ListVector intListVector = new ListVector("intList", allocator, intListField, null); + ListVector stringListVector = new ListVector("stringList", allocator, stringListField, null); + ListVector dateListVector = new ListVector("dateList", allocator, dateListField, null); + + intListVector.initializeChildrenFromFields(Arrays.asList(intField)); + stringListVector.initializeChildrenFromFields(Arrays.asList(stringField)); + dateListVector.initializeChildrenFromFields(Arrays.asList(dateField)); + + // Set up VSR + List vectors = Arrays.asList(intListVector, stringListVector, dateListVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + FieldWriter intListWriter = intListVector.getWriter(); + FieldWriter stringListWriter = stringListVector.getWriter(); + FieldWriter dateListWriter = dateListVector.getWriter(); + + // Set test data for intList + for (int i = 0; i < rowCount; i++) { + intListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + intListWriter.writeInt(j); + } + intListWriter.endList(); + } + + // Set test data for stringList + for (int i = 0; i < rowCount; i++) { + stringListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + stringListWriter.writeVarChar("string" + j); + } + stringListWriter.endList(); + } + + // Set test data for dateList + for (int i = 0; i < rowCount; i++) { + dateListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + dateListWriter.writeDateDay((int) LocalDate.now().plusDays(j).toEpochDay()); + } + dateListWriter.endList(); + } + + // Update count for the vectors + intListVector.setValueCount(rowCount); + stringListVector.setValueCount(rowCount); + dateListVector.setValueCount(rowCount); + + File dataFile = new File(TMP, "testWriteLists.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); + 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()); + 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); + assertEquals(intListVector.getObject(row), record.get("intList")); + assertEquals(dateListVector.getObject(row), record.get("dateList")); + // Handle conversion from Arrow Text type + List vectorList = stringListVector.getObject(row); + List recordList = (List) record.get("stringList"); + assertEquals(vectorList.size(), recordList.size()); + for (int i = 0; i < vectorList.size(); i++) { + assertEquals(vectorList.get(i).toString(), recordList.get(i).toString()); + } + } + } + } + } + + @Test + public void testWriteNullableLists() throws Exception { + + // Field definitions + FieldType nullListType = new FieldType(true, new ArrowType.List(), null); + FieldType nonNullListType = new FieldType(false, new ArrowType.List(), null); + + Field nullFieldType = new Field("item", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field nonNullFieldType = + new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + ListVector nullEntriesVector = + new ListVector("nullEntriesVector", allocator, nonNullListType, null); + ListVector nullListVector = new ListVector("nullListVector", allocator, nullListType, null); + ListVector nullBothVector = new ListVector("nullBothVector", allocator, nullListType, null); + + nullEntriesVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + nullListVector.initializeChildrenFromFields(Arrays.asList(nonNullFieldType)); + nullBothVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + + // Set up VSR + List vectors = Arrays.asList(nullEntriesVector, nullListVector, nullBothVector); + int rowCount = 4; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for nullEntriesVector + FieldWriter nullEntriesWriter = nullEntriesVector.getWriter(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.endList(); + + // Set test data for nullListVector + FieldWriter nullListWriter = nullListVector.getWriter(); + nullListWriter.writeNull(); + nullListWriter.setPosition(1); // writeNull() does not inc. idx() on list vector + nullListWriter.startList(); + nullListWriter.integer().writeInt(0); + nullListWriter.integer().writeInt(0); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.endList(); + + // Set test data for nullBothVector + FieldWriter nullBothWriter = nullBothVector.getWriter(); + nullBothWriter.writeNull(); + nullBothWriter.setPosition(1); + nullBothWriter.startList(); + nullBothWriter.integer().writeNull(); + nullBothWriter.integer().writeNull(); + nullBothWriter.endList(); + nullListWriter.startList(); + nullListWriter.integer().writeInt(0); + nullListWriter.integer().writeInt(0); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + + // Update count for the vectors + nullListVector.setValueCount(4); + nullEntriesVector.setValueCount(4); + nullBothVector.setValueCount(4); + + File dataFile = new File(TMP, "testWriteNullableLists.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); + 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()); + 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); + for (String list : + Arrays.asList("nullEntriesVector", "nullListVector", "nullBothVector")) { + ListVector vector = (ListVector) root.getVector(list); + Object recordField = record.get(list); + if (vector.isNull(row)) { + assertNull(recordField); + } else { + assertEquals(vector.getObject(row), recordField); + } + } + } + } + } + } + + @Test + public void testWriteFixedLists() throws Exception { + + // Field definitions + FieldType intListField = new FieldType(false, new ArrowType.FixedSizeList(5), null); + FieldType stringListField = new FieldType(false, new ArrowType.FixedSizeList(5), null); + FieldType dateListField = new FieldType(false, new ArrowType.FixedSizeList(5), null); + + Field intField = new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("item", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("item", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + FixedSizeListVector intListVector = + new FixedSizeListVector("intList", allocator, intListField, null); + FixedSizeListVector stringListVector = + new FixedSizeListVector("stringList", allocator, stringListField, null); + FixedSizeListVector dateListVector = + new FixedSizeListVector("dateList", allocator, dateListField, null); + + intListVector.initializeChildrenFromFields(Arrays.asList(intField)); + stringListVector.initializeChildrenFromFields(Arrays.asList(stringField)); + dateListVector.initializeChildrenFromFields(Arrays.asList(dateField)); + + // Set up VSR + List vectors = Arrays.asList(intListVector, stringListVector, dateListVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + FieldWriter intListWriter = intListVector.getWriter(); + FieldWriter stringListWriter = stringListVector.getWriter(); + FieldWriter dateListWriter = dateListVector.getWriter(); + + // Set test data for intList + for (int i = 0; i < rowCount; i++) { + intListWriter.startList(); + for (int j = 0; j < 5; j++) { + intListWriter.writeInt(j); + } + intListWriter.endList(); + } + + // Set test data for stringList + for (int i = 0; i < rowCount; i++) { + stringListWriter.startList(); + for (int j = 0; j < 5; j++) { + stringListWriter.writeVarChar("string" + j); + } + stringListWriter.endList(); + } + + // Set test data for dateList + for (int i = 0; i < rowCount; i++) { + dateListWriter.startList(); + for (int j = 0; j < 5; j++) { + dateListWriter.writeDateDay((int) LocalDate.now().plusDays(j).toEpochDay()); + } + dateListWriter.endList(); + } + File dataFile = new File(TMP, "testWriteFixedLists.avro"); + + // Update count for the vectors + intListVector.setValueCount(rowCount); + stringListVector.setValueCount(rowCount); + dateListVector.setValueCount(rowCount); + + // 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); + 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()); + 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); + assertEquals(intListVector.getObject(row), record.get("intList")); + assertEquals(dateListVector.getObject(row), record.get("dateList")); + // Handle conversion from Arrow Text type + List vectorList = stringListVector.getObject(row); + List recordList = (List) record.get("stringList"); + assertEquals(vectorList.size(), recordList.size()); + for (int i = 0; i < vectorList.size(); i++) { + assertEquals(vectorList.get(i).toString(), recordList.get(i).toString()); + } + } + } + } + } + + @Test + public void testWriteNullableFixedLists() throws Exception { + + // Field definitions + FieldType nullListType = new FieldType(true, new ArrowType.FixedSizeList(2), null); + FieldType nonNullListType = new FieldType(false, new ArrowType.FixedSizeList(2), null); + + Field nullFieldType = new Field("item", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field nonNullFieldType = + new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + FixedSizeListVector nullEntriesVector = + new FixedSizeListVector("nullEntriesVector", allocator, nonNullListType, null); + FixedSizeListVector nullListVector = + new FixedSizeListVector("nullListVector", allocator, nullListType, null); + FixedSizeListVector nullBothVector = + new FixedSizeListVector("nullBothVector", allocator, nullListType, null); + + nullEntriesVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + nullListVector.initializeChildrenFromFields(Arrays.asList(nonNullFieldType)); + nullBothVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + + // Set up VSR + List vectors = Arrays.asList(nullEntriesVector, nullListVector, nullBothVector); + int rowCount = 4; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for nullEntriesVector + FieldWriter nullEntriesWriter = nullEntriesVector.getWriter(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.endList(); + + // Set test data for nullListVector + FieldWriter nullListWriter = nullListVector.getWriter(); + nullListWriter.writeNull(); + nullListWriter.setPosition(1); // writeNull() does not inc. idx() on list vector + nullListWriter.startList(); + nullListWriter.integer().writeInt(123); + nullListWriter.integer().writeInt(456); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(12345); + nullEntriesWriter.integer().writeInt(67891); + nullEntriesWriter.endList(); + + // Set test data for nullBothVector + FieldWriter nullBothWriter = nullBothVector.getWriter(); + nullBothWriter.writeNull(); + nullBothWriter.setPosition(1); + nullBothWriter.startList(); + nullListWriter.integer().writeNull(); + nullListWriter.integer().writeNull(); + nullBothWriter.endList(); + nullListWriter.startList(); + nullListWriter.integer().writeInt(123); + nullListWriter.integer().writeInt(456); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + + // Update count for the vectors + nullListVector.setValueCount(4); + nullEntriesVector.setValueCount(4); + nullBothVector.setValueCount(4); + + File dataFile = new File(TMP, "testWriteNullableFixedLists.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); + 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()); + 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); + for (String list : + Arrays.asList("nullEntriesVector", "nullListVector", "nullBothVector")) { + FixedSizeListVector vector = (FixedSizeListVector) root.getVector(list); + Object recordField = record.get(list); + if (vector.isNull(row)) { + assertNull(recordField); + } else { + assertEquals(vector.getObject(row), recordField); + } + } + } + } + } + } + + @Test + public void testWriteMap() throws Exception { + + // Field definitions + FieldType intMapField = new FieldType(false, new ArrowType.Map(false), null); + FieldType stringMapField = new FieldType(false, new ArrowType.Map(false), null); + FieldType dateMapField = new FieldType(false, new ArrowType.Map(false), null); + + Field keyField = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field intField = new Field("value", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("value", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("value", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + Field intEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, intField)); + Field stringEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, stringField)); + Field dateEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, dateField)); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + MapVector intMapVector = new MapVector("intMap", allocator, intMapField, null); + MapVector stringMapVector = new MapVector("stringMap", allocator, stringMapField, null); + MapVector dateMapVector = new MapVector("dateMap", allocator, dateMapField, null); + + intMapVector.initializeChildrenFromFields(Arrays.asList(intEntryField)); + stringMapVector.initializeChildrenFromFields(Arrays.asList(stringEntryField)); + dateMapVector.initializeChildrenFromFields(Arrays.asList(dateEntryField)); + + // Set up VSR + List vectors = Arrays.asList(intMapVector, stringMapVector, dateMapVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Total number of entries that will be writen to each vector + int entryCount = 5 + 4 + 3; + + // Set test data for intList + BaseWriter.MapWriter writer = intMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + writer.startMap(); + for (int j = 0; j < 5 - i; j++) { + writer.startEntry(); + writer.key().varChar().writeVarChar("key" + j); + writer.value().integer().writeInt(j); + writer.endEntry(); + } + writer.endMap(); + } + + // Update count for data vector (map writer does not do this) + intMapVector.getDataVector().setValueCount(entryCount); + + // Set test data for stringList + BaseWriter.MapWriter stringWriter = stringMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + stringWriter.startMap(); + for (int j = 0; j < 5 - i; j++) { + stringWriter.startEntry(); + stringWriter.key().varChar().writeVarChar("key" + j); + stringWriter.value().varChar().writeVarChar("string" + j); + stringWriter.endEntry(); + } + stringWriter.endMap(); + } + + // Update count for the vectors + intMapVector.setValueCount(rowCount); + stringMapVector.setValueCount(rowCount); + dateMapVector.setValueCount(rowCount); + + // Update count for data vector (map writer does not do this) + stringMapVector.getDataVector().setValueCount(entryCount); + + // Set test data for dateList + BaseWriter.MapWriter dateWriter = dateMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + dateWriter.startMap(); + for (int j = 0; j < 5 - i; j++) { + dateWriter.startEntry(); + dateWriter.key().varChar().writeVarChar("key" + j); + dateWriter.value().dateDay().writeDateDay((int) LocalDate.now().plusDays(j).toEpochDay()); + dateWriter.endEntry(); + } + dateWriter.endMap(); + } + + // Update count for data vector (map writer does not do this) + dateMapVector.getDataVector().setValueCount(entryCount); + + File dataFile = new File(TMP, "testWriteMap.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); + 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()); + 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); + Map intMap = convertMap(intMapVector.getObject(row)); + Map stringMap = convertMap(stringMapVector.getObject(row)); + Map dateMap = convertMap(dateMapVector.getObject(row)); + compareMaps(intMap, (Map) record.get("intMap")); + compareMaps(stringMap, (Map) record.get("stringMap")); + compareMaps(dateMap, (Map) record.get("dateMap")); + } + } + } + } + + @Test + public void testWriteNullableMap() throws Exception { + + // Field definitions + FieldType nullMapType = new FieldType(true, new ArrowType.Map(false), null); + FieldType nonNullMapType = new FieldType(false, new ArrowType.Map(false), null); + + Field keyField = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field nullFieldType = new Field("value", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field nonNullFieldType = + new Field("value", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field nullEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, nullFieldType)); + Field nonNullEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, nonNullFieldType)); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + MapVector nullEntriesVector = + new MapVector("nullEntriesVector", allocator, nonNullMapType, null); + MapVector nullMapVector = new MapVector("nullMapVector", allocator, nullMapType, null); + MapVector nullBothVector = new MapVector("nullBothVector", allocator, nullMapType, null); + + nullEntriesVector.initializeChildrenFromFields(Arrays.asList(nullEntryField)); + nullMapVector.initializeChildrenFromFields(Arrays.asList(nonNullEntryField)); + nullBothVector.initializeChildrenFromFields(Arrays.asList(nullEntryField)); + + // Set up VSR + List vectors = Arrays.asList(nullEntriesVector, nullMapVector, nullBothVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for intList + BaseWriter.MapWriter writer = nullEntriesVector.getWriter(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key0"); + writer.value().integer().writeNull(); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key1"); + writer.value().integer().writeNull(); + writer.endEntry(); + writer.endMap(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key2"); + writer.value().integer().writeInt(0); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key3"); + writer.value().integer().writeInt(0); + writer.endEntry(); + writer.endMap(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + writer.endMap(); + + // Set test data for stringList + BaseWriter.MapWriter nullMapWriter = nullMapVector.getWriter(); + nullMapWriter.writeNull(); + nullMapWriter.setPosition(1); // writeNull() does not inc. idx() on map (list) vector + nullMapWriter.startMap(); + nullMapWriter.startEntry(); + nullMapWriter.key().varChar().writeVarChar("key2"); + nullMapWriter.value().integer().writeInt(0); + nullMapWriter.endEntry(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key3"); + writer.value().integer().writeInt(0); + writer.endEntry(); + nullMapWriter.endMap(); + nullMapWriter.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + nullMapWriter.endMap(); + + // Set test data for dateList + BaseWriter.MapWriter nullBothWriter = nullBothVector.getWriter(); + nullBothWriter.writeNull(); + nullBothWriter.setPosition(1); + nullBothWriter.startMap(); + nullBothWriter.startEntry(); + nullBothWriter.key().varChar().writeVarChar("key2"); + nullBothWriter.value().integer().writeNull(); + nullBothWriter.endEntry(); + nullBothWriter.startEntry(); + nullBothWriter.key().varChar().writeVarChar("key3"); + nullBothWriter.value().integer().writeNull(); + nullBothWriter.endEntry(); + nullBothWriter.endMap(); + nullBothWriter.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + nullBothWriter.endMap(); + + // Update count for the vectors + nullEntriesVector.setValueCount(3); + nullMapVector.setValueCount(3); + nullBothVector.setValueCount(3); + + File dataFile = new File(TMP, "testWriteNullableMap.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); + 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()); + 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); + Map intMap = convertMap(nullEntriesVector.getObject(row)); + Map stringMap = convertMap(nullMapVector.getObject(row)); + Map dateMap = convertMap(nullBothVector.getObject(row)); + compareMaps(intMap, (Map) record.get("nullEntriesVector")); + compareMaps(stringMap, (Map) record.get("nullMapVector")); + compareMaps(dateMap, (Map) record.get("nullBothVector")); + } + } + } + } + + private Map convertMap(List entryList) { + + if (entryList == null) { + return null; + } + + Map map = new HashMap<>(); + JsonStringArrayList structList = (JsonStringArrayList) entryList; + for (Object entry : structList) { + JsonStringHashMap structEntry = (JsonStringHashMap) entry; + String key = structEntry.get(MapVector.KEY_NAME).toString(); + Object value = structEntry.get(MapVector.VALUE_NAME); + map.put(key, value); + } + return map; + } + + private void compareMaps(Map expected, Map actual) { + if (expected == null) { + assertNull(actual); + } else { + assertEquals(expected.size(), actual.size()); + for (Object key : actual.keySet()) { + assertTrue(expected.containsKey(key.toString())); + Object actualValue = actual.get(key); + if (actualValue instanceof Utf8) { + assertEquals(expected.get(key.toString()).toString(), actualValue.toString()); + } else { + assertEquals(expected.get(key.toString()), actual.get(key)); + } + } + } + } + + @Test + public void testWriteStruct() throws Exception { + + // Field definitions + FieldType structFieldType = new FieldType(false, new ArrowType.Struct(), null); + Field intField = + new Field("intField", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("stringField", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("dateField", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + StructVector structVector = new StructVector("struct", allocator, structFieldType, null); + structVector.initializeChildrenFromFields(Arrays.asList(intField, stringField, dateField)); + + // Set up VSR + List vectors = Arrays.asList(structVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + BaseWriter.StructWriter structWriter = structVector.getWriter(); + + for (int i = 0; i < rowCount; i++) { + structWriter.start(); + structWriter.integer("intField").writeInt(i); + structWriter.varChar("stringField").writeVarChar("string" + i); + structWriter.dateDay("dateField").writeDateDay((int) LocalDate.now().toEpochDay() + i); + structWriter.end(); + } + + File dataFile = new File(TMP, "testWriteStruct.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); + for (int row = 0; row < rowCount; row++) { + producer.produce(encoder); + } + encoder.flush(); + } + + // Update count for the vector + structVector.setValueCount(rowCount); + + // Set up reading the AVRO block as a GenericRecord + Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields()); + 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); + assertNotNull(record.get("struct")); + GenericRecord structRecord = (GenericRecord) record.get("struct"); + assertEquals(row, structRecord.get("intField")); + assertEquals("string" + row, structRecord.get("stringField").toString()); + assertEquals((int) LocalDate.now().toEpochDay() + row, structRecord.get("dateField")); + } + } + } + } + + @Test + public void testWriteNullableStructs() throws Exception { + + // Field definitions + FieldType structFieldType = new FieldType(false, new ArrowType.Struct(), null); + FieldType nullableStructFieldType = new FieldType(true, new ArrowType.Struct(), null); + Field intField = + new Field("intField", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field nullableIntField = + new Field("nullableIntField", FieldType.nullable(new ArrowType.Int(32, true)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + StructVector structVector = new StructVector("struct", allocator, structFieldType, null); + StructVector nullableStructVector = + new StructVector("nullableStruct", allocator, nullableStructFieldType, null); + structVector.initializeChildrenFromFields(Arrays.asList(intField, nullableIntField)); + nullableStructVector.initializeChildrenFromFields(Arrays.asList(intField, nullableIntField)); + + // Set up VSR + List vectors = Arrays.asList(structVector, nullableStructVector); + int rowCount = 4; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for structVector + BaseWriter.StructWriter structWriter = structVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + structWriter.setPosition(i); + structWriter.start(); + structWriter.integer("intField").writeInt(i); + if (i % 2 == 0) { + structWriter.integer("nullableIntField").writeInt(i * 10); + } else { + structWriter.integer("nullableIntField").writeNull(); + } + structWriter.end(); + } + + // Set test data for nullableStructVector + BaseWriter.StructWriter nullableStructWriter = nullableStructVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + nullableStructWriter.setPosition(i); + if (i >= 2) { + nullableStructWriter.start(); + nullableStructWriter.integer("intField").writeInt(i); + if (i % 2 == 0) { + nullableStructWriter.integer("nullableIntField").writeInt(i * 10); + } else { + nullableStructWriter.integer("nullableIntField").writeNull(); + } + nullableStructWriter.end(); + } else { + nullableStructWriter.writeNull(); + } + } + + // Update count for the vector + structVector.setValueCount(rowCount); + nullableStructVector.setValueCount(rowCount); + + File dataFile = new File(TMP, "testWriteNullableStructs.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); + 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()); + 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); + if (row % 2 == 0) { + assertNotNull(record.get("struct")); + GenericRecord structRecord = (GenericRecord) record.get("struct"); + assertEquals(row, structRecord.get("intField")); + assertEquals(row * 10, structRecord.get("nullableIntField")); + } else { + assertNotNull(record.get("struct")); + GenericRecord structRecord = (GenericRecord) record.get("struct"); + assertEquals(row, structRecord.get("intField")); + assertNull(structRecord.get("nullableIntField")); + } + if (row >= 2) { + assertNotNull(record.get("nullableStruct")); + GenericRecord nullableStructRecord = (GenericRecord) record.get("nullableStruct"); + assertEquals(row, nullableStructRecord.get("intField")); + if (row % 2 == 0) { + assertEquals(row * 10, nullableStructRecord.get("nullableIntField")); + } else { + assertNull(nullableStructRecord.get("nullableIntField")); + } + } else { + assertNull(record.get("nullableStruct")); + } + } + } + } + } + + @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 new file mode 100644 index 0000000000..d5e0357a8c --- /dev/null +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroSchemaTest.java @@ -0,0 +1,1521 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +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; +import org.apache.avro.Schema; +import org.junit.jupiter.api.Test; + +public class ArrowToAvroSchemaTest { + + // Schema conversion for primitive types, nullable and non-nullable + + @Test + public void testConvertNullType() { + List fields = + Arrays.asList(new Field("nullType", FieldType.notNullable(new ArrowType.Null()), null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(1, schema.getFields().size()); + + assertEquals(Schema.Type.NULL, schema.getField("nullType").schema().getType()); + } + + @Test + public void testConvertBooleanTypes() { + List fields = + Arrays.asList( + new Field("nullableBool", FieldType.nullable(new ArrowType.Bool()), null), + new Field("nonNullableBool", FieldType.notNullable(new ArrowType.Bool()), null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(2, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableBool").schema().getType()); + assertEquals(2, schema.getField("nullableBool").schema().getTypes().size()); + assertEquals( + Schema.Type.BOOLEAN, schema.getField("nullableBool").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableBool").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.BOOLEAN, schema.getField("nonNullableBool").schema().getType()); + } + + @Test + public void testConvertIntegralTypes() { + List fields = + Arrays.asList( + new Field("nullableInt8", FieldType.nullable(new ArrowType.Int(8, true)), null), + new Field("nonNullableInt8", FieldType.notNullable(new ArrowType.Int(8, true)), null), + new Field("nullableUInt8", FieldType.nullable(new ArrowType.Int(8, false)), null), + new Field("nonNullableUInt8", FieldType.notNullable(new ArrowType.Int(8, false)), null), + new Field("nullableInt16", FieldType.nullable(new ArrowType.Int(16, true)), null), + new Field("nonNullableInt16", FieldType.notNullable(new ArrowType.Int(16, true)), null), + new Field("nullableUInt16", FieldType.nullable(new ArrowType.Int(16, false)), null), + new Field( + "nonNullableUInt16", FieldType.notNullable(new ArrowType.Int(16, false)), null), + new Field("nullableInt32", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("nonNullableInt32", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field("nullableUInt32", FieldType.nullable(new ArrowType.Int(32, false)), null), + new Field( + "nonNullableUInt32", FieldType.notNullable(new ArrowType.Int(32, false)), null), + new Field("nullableInt64", FieldType.nullable(new ArrowType.Int(64, true)), null), + new Field("nonNullableInt64", FieldType.notNullable(new ArrowType.Int(64, true)), null), + new Field("nullableUInt64", FieldType.nullable(new ArrowType.Int(64, false)), null), + new Field( + "nonNullableUInt64", FieldType.notNullable(new ArrowType.Int(64, false)), null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(16, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableInt8").schema().getType()); + assertEquals(2, schema.getField("nullableInt8").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("nullableInt8").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableInt8").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.INT, schema.getField("nonNullableInt8").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableUInt8").schema().getType()); + assertEquals(2, schema.getField("nullableUInt8").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("nullableUInt8").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableUInt8").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.INT, schema.getField("nonNullableUInt8").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableInt16").schema().getType()); + assertEquals(2, schema.getField("nullableInt16").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("nullableInt16").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableInt16").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.INT, schema.getField("nonNullableInt16").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableUInt16").schema().getType()); + assertEquals(2, schema.getField("nullableUInt16").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("nullableUInt16").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableUInt16").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.INT, schema.getField("nonNullableUInt16").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableInt32").schema().getType()); + assertEquals(2, schema.getField("nullableInt32").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("nullableInt32").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableInt32").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.INT, schema.getField("nonNullableInt32").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableUInt32").schema().getType()); + assertEquals(2, schema.getField("nullableUInt32").schema().getTypes().size()); + assertEquals( + Schema.Type.LONG, schema.getField("nullableUInt32").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableUInt32").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.LONG, schema.getField("nonNullableUInt32").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableInt64").schema().getType()); + assertEquals(2, schema.getField("nullableInt64").schema().getTypes().size()); + assertEquals( + Schema.Type.LONG, schema.getField("nullableInt64").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableInt64").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.LONG, schema.getField("nonNullableInt64").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableUInt64").schema().getType()); + assertEquals(2, schema.getField("nullableUInt64").schema().getTypes().size()); + assertEquals( + Schema.Type.LONG, schema.getField("nullableUInt64").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableUInt64").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.LONG, schema.getField("nonNullableUInt64").schema().getType()); + } + + @Test + public void testConvertFloatingPointTypes() { + List fields = + Arrays.asList( + new Field( + "nullableFloat16", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.HALF)), + null), + new Field( + "nonNullableFloat16", + FieldType.notNullable(new ArrowType.FloatingPoint(FloatingPointPrecision.HALF)), + null), + new Field( + "nullableFloat32", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nonNullableFloat32", + FieldType.notNullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nullableFloat64", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "nonNullableFloat64", + FieldType.notNullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(6, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableFloat16").schema().getType()); + assertEquals(2, schema.getField("nullableFloat16").schema().getTypes().size()); + assertEquals( + Schema.Type.FLOAT, schema.getField("nullableFloat16").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableFloat16").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.FLOAT, schema.getField("nonNullableFloat16").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableFloat32").schema().getType()); + assertEquals(2, schema.getField("nullableFloat32").schema().getTypes().size()); + assertEquals( + Schema.Type.FLOAT, schema.getField("nullableFloat32").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableFloat32").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.FLOAT, schema.getField("nonNullableFloat32").schema().getType()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableFloat64").schema().getType()); + assertEquals(2, schema.getField("nullableFloat64").schema().getTypes().size()); + assertEquals( + Schema.Type.DOUBLE, + schema.getField("nullableFloat64").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableFloat64").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.DOUBLE, schema.getField("nonNullableFloat64").schema().getType()); + } + + @Test + public void testConvertStringTypes() { + List fields = + Arrays.asList( + new Field("nullableUtf8", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("nonNullableUtf8", FieldType.notNullable(new ArrowType.Utf8()), null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(2, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableUtf8").schema().getType()); + assertEquals(2, schema.getField("nullableUtf8").schema().getTypes().size()); + assertEquals( + Schema.Type.STRING, schema.getField("nullableUtf8").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableUtf8").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.STRING, schema.getField("nonNullableUtf8").schema().getType()); + } + + @Test + public void testConvertBinaryTypes() { + List fields = + Arrays.asList( + new Field("nullableBinary", FieldType.nullable(new ArrowType.Binary()), null), + new Field("nonNullableBinary", FieldType.notNullable(new ArrowType.Binary()), null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(2, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableBinary").schema().getType()); + assertEquals(2, schema.getField("nullableBinary").schema().getTypes().size()); + assertEquals( + Schema.Type.BYTES, schema.getField("nullableBinary").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableBinary").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.BYTES, schema.getField("nonNullableBinary").schema().getType()); + } + + @Test + public void testConvertFixedSizeBinaryTypes() { + List fields = + Arrays.asList( + new Field( + "nullableFixedSizeBinary", + FieldType.nullable(new ArrowType.FixedSizeBinary(10)), + null), + new Field( + "nonNullableFixedSizeBinary", + FieldType.notNullable(new ArrowType.FixedSizeBinary(10)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(2, schema.getFields().size()); + + assertEquals(Schema.Type.UNION, schema.getField("nullableFixedSizeBinary").schema().getType()); + assertEquals(2, schema.getField("nullableFixedSizeBinary").schema().getTypes().size()); + Schema nullableFixedSizeBinarySchema = + schema.getField("nullableFixedSizeBinary").schema().getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableFixedSizeBinarySchema.getType()); + assertEquals(10, nullableFixedSizeBinarySchema.getFixedSize()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableFixedSizeBinary").schema().getTypes().get(1).getType()); + Schema nonNullableFixedSizeBinarySchema = + schema.getField("nullableFixedSizeBinary").schema().getTypes().get(0); + assertEquals(Schema.Type.FIXED, nonNullableFixedSizeBinarySchema.getType()); + assertEquals(10, nonNullableFixedSizeBinarySchema.getFixedSize()); + } + + // Schema conversion for logical types, nullable and non-nullable + + @Test + public void testConvertDecimalTypes() { + List fields = + Arrays.asList( + new Field( + "nullableDecimal128", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "nonNullableDecimal1281", + FieldType.notNullable(new ArrowType.Decimal(10, 2, 128)), + null), + new Field( + "nonNullableDecimal1282", + FieldType.notNullable(new ArrowType.Decimal(15, 5, 128)), + null), + new Field( + "nonNullableDecimal1283", + FieldType.notNullable(new ArrowType.Decimal(20, 10, 128)), + null), + new Field( + "nullableDecimal256", FieldType.nullable(new ArrowType.Decimal(55, 15, 256)), null), + new Field( + "nonNullableDecimal2561", + FieldType.notNullable(new ArrowType.Decimal(55, 25, 256)), + null), + new Field( + "nonNullableDecimal2562", + FieldType.notNullable(new ArrowType.Decimal(25, 8, 256)), + null), + new Field( + "nonNullableDecimal2563", + FieldType.notNullable(new ArrowType.Decimal(60, 50, 256)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(8, schema.getFields().size()); + + // Assertions for nullableDecimal128 + assertEquals(Schema.Type.UNION, schema.getField("nullableDecimal128").schema().getType()); + assertEquals(2, schema.getField("nullableDecimal128").schema().getTypes().size()); + Schema nullableDecimal128Schema = + schema.getField("nullableDecimal128").schema().getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableDecimal128Schema.getType()); + assertEquals(16, nullableDecimal128Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(10, 2), nullableDecimal128Schema.getLogicalType()); + assertEquals(10, nullableDecimal128Schema.getObjectProp("precision")); + assertEquals(2, nullableDecimal128Schema.getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableDecimal128").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableDecimal1281 + Schema nonNullableDecimal1281Schema = schema.getField("nonNullableDecimal1281").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal1281Schema.getType()); + assertEquals(16, nonNullableDecimal1281Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(10, 2), nonNullableDecimal1281Schema.getLogicalType()); + assertEquals(10, nonNullableDecimal1281Schema.getObjectProp("precision")); + assertEquals(2, nonNullableDecimal1281Schema.getObjectProp("scale")); + + // Assertions for nonNullableDecimal1282 + Schema nonNullableDecimal1282Schema = schema.getField("nonNullableDecimal1282").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal1282Schema.getType()); + assertEquals(16, nonNullableDecimal1282Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(15, 5), nonNullableDecimal1282Schema.getLogicalType()); + assertEquals(15, nonNullableDecimal1282Schema.getObjectProp("precision")); + assertEquals(5, nonNullableDecimal1282Schema.getObjectProp("scale")); + + // Assertions for nonNullableDecimal1283 + Schema nonNullableDecimal1283Schema = schema.getField("nonNullableDecimal1283").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal1283Schema.getType()); + assertEquals(16, nonNullableDecimal1283Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(20, 10), nonNullableDecimal1283Schema.getLogicalType()); + assertEquals(20, nonNullableDecimal1283Schema.getObjectProp("precision")); + assertEquals(10, nonNullableDecimal1283Schema.getObjectProp("scale")); + + // Assertions for nullableDecimal256 + assertEquals(Schema.Type.UNION, schema.getField("nullableDecimal256").schema().getType()); + assertEquals(2, schema.getField("nullableDecimal256").schema().getTypes().size()); + Schema nullableDecimal256Schema = + schema.getField("nullableDecimal256").schema().getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableDecimal256Schema.getType()); + assertEquals(32, nullableDecimal256Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(55, 15), nullableDecimal256Schema.getLogicalType()); + assertEquals(55, nullableDecimal256Schema.getObjectProp("precision")); + assertEquals(15, nullableDecimal256Schema.getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableDecimal256").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableDecimal2561 + Schema nonNullableDecimal2561Schema = schema.getField("nonNullableDecimal2561").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal2561Schema.getType()); + assertEquals(32, nonNullableDecimal2561Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(55, 25), nonNullableDecimal2561Schema.getLogicalType()); + assertEquals(55, nonNullableDecimal2561Schema.getObjectProp("precision")); + assertEquals(25, nonNullableDecimal2561Schema.getObjectProp("scale")); + + // Assertions for nonNullableDecimal2562 + Schema nonNullableDecimal2562Schema = schema.getField("nonNullableDecimal2562").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal2562Schema.getType()); + assertEquals(32, nonNullableDecimal2562Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(25, 8), nonNullableDecimal2562Schema.getLogicalType()); + assertEquals(25, nonNullableDecimal2562Schema.getObjectProp("precision")); + assertEquals(8, nonNullableDecimal2562Schema.getObjectProp("scale")); + + // Assertions for nonNullableDecimal2563 + Schema nonNullableDecimal2563Schema = schema.getField("nonNullableDecimal2563").schema(); + assertEquals(Schema.Type.FIXED, nonNullableDecimal2563Schema.getType()); + assertEquals(32, nonNullableDecimal2563Schema.getFixedSize()); + assertEquals(LogicalTypes.decimal(60, 50), nonNullableDecimal2563Schema.getLogicalType()); + assertEquals(60, nonNullableDecimal2563Schema.getObjectProp("precision")); + assertEquals(50, nonNullableDecimal2563Schema.getObjectProp("scale")); + } + + @Test + public void testConvertDateTypes() { + List fields = + Arrays.asList( + new Field( + "nullableDateDay", FieldType.nullable(new ArrowType.Date(DateUnit.DAY)), null), + new Field( + "nonNullableDateDay", + FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), + null), + new Field( + "nullableDateMilli", + FieldType.nullable(new ArrowType.Date(DateUnit.MILLISECOND)), + null), + new Field( + "nonNullableDateMilli", + FieldType.notNullable(new ArrowType.Date(DateUnit.MILLISECOND)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(4, schema.getFields().size()); + + // Assertions for nullableDateDay + assertEquals(Schema.Type.UNION, schema.getField("nullableDateDay").schema().getType()); + assertEquals(2, schema.getField("nullableDateDay").schema().getTypes().size()); + Schema nullableDateDaySchema = schema.getField("nullableDateDay").schema().getTypes().get(0); + assertEquals(Schema.Type.INT, nullableDateDaySchema.getType()); + assertEquals(LogicalTypes.date(), nullableDateDaySchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableDateDay").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableDateDay + Schema nonNullableDateDaySchema = schema.getField("nonNullableDateDay").schema(); + assertEquals(Schema.Type.INT, nonNullableDateDaySchema.getType()); + assertEquals(LogicalTypes.date(), nonNullableDateDaySchema.getLogicalType()); + + // Assertions for nullableDateMilli + assertEquals(Schema.Type.UNION, schema.getField("nullableDateMilli").schema().getType()); + assertEquals(2, schema.getField("nullableDateMilli").schema().getTypes().size()); + Schema nullableDateMilliSchema = + schema.getField("nullableDateMilli").schema().getTypes().get(0); + assertEquals(Schema.Type.INT, nullableDateMilliSchema.getType()); + assertEquals(LogicalTypes.date(), nullableDateMilliSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableDateMilli").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableDateMilli + Schema nonNullableDateMilliSchema = schema.getField("nonNullableDateMilli").schema(); + assertEquals(Schema.Type.INT, nonNullableDateMilliSchema.getType()); + assertEquals(LogicalTypes.date(), nonNullableDateMilliSchema.getLogicalType()); + } + + @Test + public void testConvertTimeTypes() { + List fields = + Arrays.asList( + new Field( + "nullableTimeSec", + FieldType.nullable(new ArrowType.Time(TimeUnit.SECOND, 32)), + null), + new Field( + "nonNullableTimeSec", + FieldType.notNullable(new ArrowType.Time(TimeUnit.SECOND, 32)), + null), + new Field( + "nullableTimeMillis", + FieldType.nullable(new ArrowType.Time(TimeUnit.MILLISECOND, 32)), + null), + new Field( + "nonNullableTimeMillis", + FieldType.notNullable(new ArrowType.Time(TimeUnit.MILLISECOND, 32)), + null), + new Field( + "nullableTimeMicros", + FieldType.nullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), + null), + new Field( + "nonNullableTimeMicros", + FieldType.notNullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), + null), + new Field( + "nullableTimeNanos", + FieldType.nullable(new ArrowType.Time(TimeUnit.NANOSECOND, 64)), + null), + new Field( + "nonNullableTimeNanos", + FieldType.notNullable(new ArrowType.Time(TimeUnit.NANOSECOND, 64)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(8, schema.getFields().size()); + + // Assertions for nullableTimeSec + assertEquals(Schema.Type.UNION, schema.getField("nullableTimeSec").schema().getType()); + assertEquals(2, schema.getField("nullableTimeSec").schema().getTypes().size()); + Schema nullableTimeSecSchema = schema.getField("nullableTimeSec").schema().getTypes().get(0); + assertEquals(Schema.Type.INT, nullableTimeSecSchema.getType()); + assertEquals(LogicalTypes.timeMillis(), nullableTimeSecSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableTimeSec").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimeSec + Schema nonNullableTimeSecSchema = schema.getField("nonNullableTimeSec").schema(); + assertEquals(Schema.Type.INT, nonNullableTimeSecSchema.getType()); + assertEquals(LogicalTypes.timeMillis(), nonNullableTimeSecSchema.getLogicalType()); + + // Assertions for nullableTimeMillis + assertEquals(Schema.Type.UNION, schema.getField("nullableTimeMillis").schema().getType()); + assertEquals(2, schema.getField("nullableTimeMillis").schema().getTypes().size()); + Schema nullableTimeMillisSchema = + schema.getField("nullableTimeMillis").schema().getTypes().get(0); + assertEquals(Schema.Type.INT, nullableTimeMillisSchema.getType()); + assertEquals(LogicalTypes.timeMillis(), nullableTimeMillisSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimeMillis").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimeMillis + Schema nonNullableTimeMillisSchema = schema.getField("nonNullableTimeMillis").schema(); + assertEquals(Schema.Type.INT, nonNullableTimeMillisSchema.getType()); + assertEquals(LogicalTypes.timeMillis(), nonNullableTimeMillisSchema.getLogicalType()); + + // Assertions for nullableTimeMicros + assertEquals(Schema.Type.UNION, schema.getField("nullableTimeMicros").schema().getType()); + assertEquals(2, schema.getField("nullableTimeMicros").schema().getTypes().size()); + Schema nullableTimeMicrosSchema = + schema.getField("nullableTimeMicros").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimeMicrosSchema.getType()); + assertEquals(LogicalTypes.timeMicros(), nullableTimeMicrosSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimeMicros").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimeMicros + Schema nonNullableTimeMicrosSchema = schema.getField("nonNullableTimeMicros").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimeMicrosSchema.getType()); + assertEquals(LogicalTypes.timeMicros(), nonNullableTimeMicrosSchema.getLogicalType()); + + // Assertions for nullableTimeNanos + assertEquals(Schema.Type.UNION, schema.getField("nullableTimeNanos").schema().getType()); + assertEquals(2, schema.getField("nullableTimeNanos").schema().getTypes().size()); + Schema nullableTimeNanosSchema = + schema.getField("nullableTimeNanos").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimeNanosSchema.getType()); + assertEquals(LogicalTypes.timeMicros(), nullableTimeNanosSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimeNanos").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimeNanos + Schema nonNullableTimeNanosSchema = schema.getField("nonNullableTimeNanos").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimeNanosSchema.getType()); + assertEquals(LogicalTypes.timeMicros(), nonNullableTimeNanosSchema.getLogicalType()); + } + + @Test + public void testConvertZoneAwareTimestampTypes() { + List fields = + Arrays.asList( + new Field( + "nullableTimestampSecTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.SECOND, "UTC")), + null), + new Field( + "nonNullableTimestampSecTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.SECOND, "UTC")), + null), + new Field( + "nullableTimestampMillisTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null), + new Field( + "nonNullableTimestampMillisTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null), + new Field( + "nullableTimestampMicrosTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null), + new Field( + "nonNullableTimestampMicrosTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null), + new Field( + "nullableTimestampNanosTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), + null), + new Field( + "nonNullableTimestampNanosTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(8, schema.getFields().size()); + + // Assertions for nullableTimestampSecTz + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampSecTz").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampSecTz").schema().getTypes().size()); + Schema nullableTimestampSecTzSchema = + schema.getField("nullableTimestampSecTz").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampSecTzSchema.getType()); + assertEquals(LogicalTypes.timestampMillis(), nullableTimestampSecTzSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampSecTz").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampSecTz + Schema nonNullableTimestampSecTzSchema = schema.getField("nonNullableTimestampSecTz").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampSecTzSchema.getType()); + assertEquals(LogicalTypes.timestampMillis(), nonNullableTimestampSecTzSchema.getLogicalType()); + + // Assertions for nullableTimestampMillisTz + assertEquals( + Schema.Type.UNION, schema.getField("nullableTimestampMillisTz").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampMillisTz").schema().getTypes().size()); + Schema nullableTimestampMillisTzSchema = + schema.getField("nullableTimestampMillisTz").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampMillisTzSchema.getType()); + assertEquals(LogicalTypes.timestampMillis(), nullableTimestampMillisTzSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampMillisTz").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampMillisTz + Schema nonNullableTimestampMillisTzSchema = + schema.getField("nonNullableTimestampMillisTz").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampMillisTzSchema.getType()); + assertEquals( + LogicalTypes.timestampMillis(), nonNullableTimestampMillisTzSchema.getLogicalType()); + + // Assertions for nullableTimestampMicrosTz + assertEquals( + Schema.Type.UNION, schema.getField("nullableTimestampMicrosTz").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampMicrosTz").schema().getTypes().size()); + Schema nullableTimestampMicrosTzSchema = + schema.getField("nullableTimestampMicrosTz").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampMicrosTzSchema.getType()); + assertEquals(LogicalTypes.timestampMicros(), nullableTimestampMicrosTzSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampMicrosTz").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampMicrosTz + Schema nonNullableTimestampMicrosTzSchema = + schema.getField("nonNullableTimestampMicrosTz").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampMicrosTzSchema.getType()); + assertEquals( + LogicalTypes.timestampMicros(), nonNullableTimestampMicrosTzSchema.getLogicalType()); + + // Assertions for nullableTimestampNanosTz + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampNanosTz").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampNanosTz").schema().getTypes().size()); + Schema nullableTimestampNanosTzSchema = + schema.getField("nullableTimestampNanosTz").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampNanosTzSchema.getType()); + assertEquals(LogicalTypes.timestampNanos(), nullableTimestampNanosTzSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampNanosTz").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampNanosTz + Schema nonNullableTimestampNanosTzSchema = + schema.getField("nonNullableTimestampNanosTz").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampNanosTzSchema.getType()); + assertEquals(LogicalTypes.timestampNanos(), nonNullableTimestampNanosTzSchema.getLogicalType()); + } + + @Test + public void testConvertLocalTimestampTypes() { + List fields = + Arrays.asList( + new Field( + "nullableTimestampSec", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.SECOND, null)), + null), + new Field( + "nonNullableTimestampSec", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.SECOND, null)), + null), + new Field( + "nullableTimestampMillis", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + null), + new Field( + "nonNullableTimestampMillis", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + null), + new Field( + "nullableTimestampMicros", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), + null), + new Field( + "nonNullableTimestampMicros", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), + null), + new Field( + "nullableTimestampNanos", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), + null), + new Field( + "nonNullableTimestampNanos", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), + null)); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(8, schema.getFields().size()); + + // Assertions for nullableTimestampSec + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampSec").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampSec").schema().getTypes().size()); + Schema nullableTimestampSecSchema = + schema.getField("nullableTimestampSec").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampSecSchema.getType()); + assertEquals(LogicalTypes.localTimestampMillis(), nullableTimestampSecSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampSec").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampSec + Schema nonNullableTimestampSecSchema = schema.getField("nonNullableTimestampSec").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampSecSchema.getType()); + assertEquals( + LogicalTypes.localTimestampMillis(), nonNullableTimestampSecSchema.getLogicalType()); + + // Assertions for nullableTimestampMillis + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampMillis").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampMillis").schema().getTypes().size()); + Schema nullableTimestampMillisSchema = + schema.getField("nullableTimestampMillis").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampMillisSchema.getType()); + assertEquals( + LogicalTypes.localTimestampMillis(), nullableTimestampMillisSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampMillis").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampMillis + Schema nonNullableTimestampMillisSchema = + schema.getField("nonNullableTimestampMillis").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampMillisSchema.getType()); + assertEquals( + LogicalTypes.localTimestampMillis(), nonNullableTimestampMillisSchema.getLogicalType()); + + // Assertions for nullableTimestampMicros + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampMicros").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampMicros").schema().getTypes().size()); + Schema nullableTimestampMicrosSchema = + schema.getField("nullableTimestampMicros").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampMicrosSchema.getType()); + assertEquals( + LogicalTypes.localTimestampMicros(), nullableTimestampMicrosSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampMicros").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampMicros + Schema nonNullableTimestampMicrosSchema = + schema.getField("nonNullableTimestampMicros").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampMicrosSchema.getType()); + assertEquals( + LogicalTypes.localTimestampMicros(), nonNullableTimestampMicrosSchema.getLogicalType()); + + // Assertions for nullableTimestampNanos + assertEquals(Schema.Type.UNION, schema.getField("nullableTimestampNanos").schema().getType()); + assertEquals(2, schema.getField("nullableTimestampNanos").schema().getTypes().size()); + Schema nullableTimestampNanosSchema = + schema.getField("nullableTimestampNanos").schema().getTypes().get(0); + assertEquals(Schema.Type.LONG, nullableTimestampNanosSchema.getType()); + assertEquals(LogicalTypes.localTimestampNanos(), nullableTimestampNanosSchema.getLogicalType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableTimestampNanos").schema().getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampNanos + Schema nonNullableTimestampNanosSchema = schema.getField("nonNullableTimestampNanos").schema(); + assertEquals(Schema.Type.LONG, nonNullableTimestampNanosSchema.getType()); + assertEquals( + LogicalTypes.localTimestampNanos(), nonNullableTimestampNanosSchema.getLogicalType()); + } + + // Schema conversion for complex types, where the contents are primitive and logical types + + @Test + public void testConvertListTypes() { + List fields = + Arrays.asList( + new Field( + "nullableIntList", + FieldType.nullable(new ArrowType.List()), + Arrays.asList( + new Field("item", FieldType.nullable(new ArrowType.Int(32, true)), null))), + new Field( + "nullableDoubleList", + FieldType.nullable(new ArrowType.List()), + Arrays.asList( + new Field( + "item", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null))), + new Field( + "nonNullableDecimalList", + FieldType.notNullable(new ArrowType.List()), + Arrays.asList( + new Field( + "item", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null))), + new Field( + "nonNullableTimestampList", + FieldType.notNullable(new ArrowType.List()), + Arrays.asList( + new Field( + "item", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(4, schema.getFields().size()); + + // Assertions for nullableIntList + assertEquals(Schema.Type.UNION, schema.getField("nullableIntList").schema().getType()); + assertEquals(2, schema.getField("nullableIntList").schema().getTypes().size()); + assertEquals( + Schema.Type.ARRAY, schema.getField("nullableIntList").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableIntList").schema().getTypes().get(1).getType()); + Schema nullableIntListItemSchema = + schema.getField("nullableIntList").schema().getTypes().get(0).getElementType(); + assertEquals(Schema.Type.UNION, nullableIntListItemSchema.getType()); + assertEquals(2, nullableIntListItemSchema.getTypes().size()); + assertEquals(Schema.Type.INT, nullableIntListItemSchema.getTypes().get(0).getType()); + assertEquals(Schema.Type.NULL, nullableIntListItemSchema.getTypes().get(1).getType()); + + // Assertions for nullableDoubleList + assertEquals(Schema.Type.UNION, schema.getField("nullableDoubleList").schema().getType()); + assertEquals(2, schema.getField("nullableDoubleList").schema().getTypes().size()); + assertEquals( + Schema.Type.ARRAY, + schema.getField("nullableDoubleList").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableDoubleList").schema().getTypes().get(1).getType()); + Schema nullableDoubleListItemSchema = + schema.getField("nullableDoubleList").schema().getTypes().get(0).getElementType(); + assertEquals(Schema.Type.DOUBLE, nullableDoubleListItemSchema.getType()); + + // Assertions for nonNullableDecimalList + assertEquals(Schema.Type.ARRAY, schema.getField("nonNullableDecimalList").schema().getType()); + Schema nonNullableDecimalListItemSchema = + schema.getField("nonNullableDecimalList").schema().getElementType(); + assertEquals(Schema.Type.UNION, nonNullableDecimalListItemSchema.getType()); + assertEquals(2, nonNullableDecimalListItemSchema.getTypes().size()); + Schema nullableDecimalSchema = nonNullableDecimalListItemSchema.getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableDecimalSchema.getType()); + assertEquals(16, nullableDecimalSchema.getFixedSize()); + assertEquals(LogicalTypes.decimal(10, 2), nullableDecimalSchema.getLogicalType()); + assertEquals(10, nullableDecimalSchema.getObjectProp("precision")); + assertEquals(2, nullableDecimalSchema.getObjectProp("scale")); + assertEquals(Schema.Type.NULL, nonNullableDecimalListItemSchema.getTypes().get(1).getType()); + + // Assertions for nonNullableTimestampList + assertEquals(Schema.Type.ARRAY, schema.getField("nonNullableTimestampList").schema().getType()); + Schema nonNullableTimestampListItemSchema = + schema.getField("nonNullableTimestampList").schema().getElementType(); + assertEquals(Schema.Type.LONG, nonNullableTimestampListItemSchema.getType()); + assertEquals( + LogicalTypes.timestampMillis(), nonNullableTimestampListItemSchema.getLogicalType()); + } + + @Test + public void testConvertFixedSizeListTypes() { + List fields = + Arrays.asList( + new Field( + "nullableFixedSizeIntList", + FieldType.nullable(new ArrowType.FixedSizeList(3)), + Arrays.asList( + new Field("item", FieldType.nullable(new ArrowType.Int(32, true)), null))), + new Field( + "nullableFixedSizeDoubleList", + FieldType.nullable(new ArrowType.FixedSizeList(3)), + Arrays.asList( + new Field( + "item", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null))), + new Field( + "nonNullableFixedSizeDecimalList", + FieldType.notNullable(new ArrowType.FixedSizeList(3)), + Arrays.asList( + new Field( + "item", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null))), + new Field( + "nonNullableFixedSizeTimestampList", + FieldType.notNullable(new ArrowType.FixedSizeList(3)), + Arrays.asList( + new Field( + "item", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(4, schema.getFields().size()); + + // Assertions for nullableFixedSizeIntList + assertEquals(Schema.Type.UNION, schema.getField("nullableFixedSizeIntList").schema().getType()); + assertEquals(2, schema.getField("nullableFixedSizeIntList").schema().getTypes().size()); + assertEquals( + Schema.Type.ARRAY, + schema.getField("nullableFixedSizeIntList").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableFixedSizeIntList").schema().getTypes().get(1).getType()); + Schema nullableFixedSizeIntListItemSchema = + schema.getField("nullableFixedSizeIntList").schema().getTypes().get(0).getElementType(); + assertEquals(Schema.Type.UNION, nullableFixedSizeIntListItemSchema.getType()); + assertEquals(2, nullableFixedSizeIntListItemSchema.getTypes().size()); + assertEquals(Schema.Type.INT, nullableFixedSizeIntListItemSchema.getTypes().get(0).getType()); + assertEquals(Schema.Type.NULL, nullableFixedSizeIntListItemSchema.getTypes().get(1).getType()); + + // Assertions for nullableFixedSizeDoubleList + assertEquals( + Schema.Type.UNION, schema.getField("nullableFixedSizeDoubleList").schema().getType()); + assertEquals(2, schema.getField("nullableFixedSizeDoubleList").schema().getTypes().size()); + assertEquals( + Schema.Type.ARRAY, + schema.getField("nullableFixedSizeDoubleList").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableFixedSizeDoubleList").schema().getTypes().get(1).getType()); + Schema nullableFixedSizeDoubleListItemSchema = + schema.getField("nullableFixedSizeDoubleList").schema().getTypes().get(0).getElementType(); + assertEquals(Schema.Type.DOUBLE, nullableFixedSizeDoubleListItemSchema.getType()); + + // Assertions for nonNullableFixedSizeDecimalList + assertEquals( + Schema.Type.ARRAY, schema.getField("nonNullableFixedSizeDecimalList").schema().getType()); + Schema nonNullableFixedSizeDecimalListItemSchema = + schema.getField("nonNullableFixedSizeDecimalList").schema().getElementType(); + assertEquals(Schema.Type.UNION, nonNullableFixedSizeDecimalListItemSchema.getType()); + assertEquals(2, nonNullableFixedSizeDecimalListItemSchema.getTypes().size()); + Schema nullableDecimalSchema = nonNullableFixedSizeDecimalListItemSchema.getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableDecimalSchema.getType()); + assertEquals(16, nullableDecimalSchema.getFixedSize()); + assertEquals(LogicalTypes.decimal(10, 2), nullableDecimalSchema.getLogicalType()); + assertEquals(10, nullableDecimalSchema.getObjectProp("precision")); + assertEquals(2, nullableDecimalSchema.getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, nonNullableFixedSizeDecimalListItemSchema.getTypes().get(1).getType()); + + // Assertions for nonNullableFixedSizeTimestampList + assertEquals( + Schema.Type.ARRAY, schema.getField("nonNullableFixedSizeTimestampList").schema().getType()); + Schema nonNullableFixedSizeTimestampListItemSchema = + schema.getField("nonNullableFixedSizeTimestampList").schema().getElementType(); + assertEquals(Schema.Type.LONG, nonNullableFixedSizeTimestampListItemSchema.getType()); + assertEquals( + LogicalTypes.timestampMillis(), + nonNullableFixedSizeTimestampListItemSchema.getLogicalType()); + } + + @Test + public void testConvertMapTypes() { + List fields = + Arrays.asList( + new Field( + "nullableMapWithNullableInt", + FieldType.nullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", FieldType.nullable(new ArrowType.Int(32, true)), null))))), + new Field( + "nullableMapWithNonNullableDouble", + FieldType.nullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null))))), + new Field( + "nonNullableMapWithNullableDecimal", + FieldType.notNullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), + null))))), + new Field( + "nonNullableMapWithNonNullableTimestamp", + FieldType.notNullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.notNullable( + new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))))); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(4, schema.getFields().size()); + + // Assertions for nullableMapWithNullableInt + assertEquals( + Schema.Type.UNION, schema.getField("nullableMapWithNullableInt").schema().getType()); + assertEquals(2, schema.getField("nullableMapWithNullableInt").schema().getTypes().size()); + assertEquals( + Schema.Type.MAP, + schema.getField("nullableMapWithNullableInt").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableMapWithNullableInt").schema().getTypes().get(1).getType()); + Schema nullableMapWithNullableIntValueSchema = + schema.getField("nullableMapWithNullableInt").schema().getTypes().get(0).getValueType(); + assertEquals(Schema.Type.UNION, nullableMapWithNullableIntValueSchema.getType()); + assertEquals(2, nullableMapWithNullableIntValueSchema.getTypes().size()); + assertEquals( + Schema.Type.INT, nullableMapWithNullableIntValueSchema.getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, nullableMapWithNullableIntValueSchema.getTypes().get(1).getType()); + + // Assertions for nullableMapWithNonNullableDouble + assertEquals( + Schema.Type.UNION, schema.getField("nullableMapWithNonNullableDouble").schema().getType()); + assertEquals(2, schema.getField("nullableMapWithNonNullableDouble").schema().getTypes().size()); + assertEquals( + Schema.Type.MAP, + schema.getField("nullableMapWithNonNullableDouble").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableMapWithNonNullableDouble").schema().getTypes().get(1).getType()); + Schema nullableMapWithNonNullableDoubleValueSchema = + schema + .getField("nullableMapWithNonNullableDouble") + .schema() + .getTypes() + .get(0) + .getValueType(); + assertEquals(Schema.Type.DOUBLE, nullableMapWithNonNullableDoubleValueSchema.getType()); + + // Assertions for nonNullableMapWithNullableDecimal + assertEquals( + Schema.Type.MAP, schema.getField("nonNullableMapWithNullableDecimal").schema().getType()); + Schema nonNullableMapWithNullableDecimalValueSchema = + schema.getField("nonNullableMapWithNullableDecimal").schema().getValueType(); + assertEquals(Schema.Type.UNION, nonNullableMapWithNullableDecimalValueSchema.getType()); + assertEquals(2, nonNullableMapWithNullableDecimalValueSchema.getTypes().size()); + Schema nullableDecimalSchema = nonNullableMapWithNullableDecimalValueSchema.getTypes().get(0); + assertEquals(Schema.Type.FIXED, nullableDecimalSchema.getType()); + assertEquals(16, nullableDecimalSchema.getFixedSize()); + assertEquals(LogicalTypes.decimal(10, 2), nullableDecimalSchema.getLogicalType()); + assertEquals(10, nullableDecimalSchema.getObjectProp("precision")); + assertEquals(2, nullableDecimalSchema.getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, nonNullableMapWithNullableDecimalValueSchema.getTypes().get(1).getType()); + + // Assertions for nonNullableMapWithNonNullableTimestamp + assertEquals( + Schema.Type.MAP, + schema.getField("nonNullableMapWithNonNullableTimestamp").schema().getType()); + Schema nonNullableMapWithNonNullableTimestampValueSchema = + schema.getField("nonNullableMapWithNonNullableTimestamp").schema().getValueType(); + assertEquals(Schema.Type.LONG, nonNullableMapWithNonNullableTimestampValueSchema.getType()); + assertEquals( + LogicalTypes.timestampMillis(), + nonNullableMapWithNonNullableTimestampValueSchema.getLogicalType()); + } + + @Test + public void testConvertRecordTypes() { + List fields = + Arrays.asList( + new Field( + "nullableRecord", + FieldType.nullable(new ArrowType.Struct()), + Arrays.asList( + new Field("field1", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "field2", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "field3", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "field4", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null))), + new Field( + "nonNullableRecord", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("field1", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "field2", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "field3", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "field4", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(2, schema.getFields().size()); + + // Assertions for nullableRecord + assertEquals(Schema.Type.UNION, schema.getField("nullableRecord").schema().getType()); + assertEquals(2, schema.getField("nullableRecord").schema().getTypes().size()); + assertEquals( + Schema.Type.RECORD, schema.getField("nullableRecord").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, schema.getField("nullableRecord").schema().getTypes().get(1).getType()); + Schema nullableRecordSchema = schema.getField("nullableRecord").schema().getTypes().get(0); + assertEquals(4, nullableRecordSchema.getFields().size()); + assertEquals( + Schema.Type.INT, + nullableRecordSchema.getField("field1").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + nullableRecordSchema.getField("field1").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.DOUBLE, nullableRecordSchema.getField("field2").schema().getType()); + assertEquals( + Schema.Type.FIXED, + nullableRecordSchema.getField("field3").schema().getTypes().get(0).getType()); + assertEquals( + 16, nullableRecordSchema.getField("field3").schema().getTypes().get(0).getFixedSize()); + assertEquals( + LogicalTypes.decimal(10, 2), + nullableRecordSchema.getField("field3").schema().getTypes().get(0).getLogicalType()); + assertEquals( + 10, + nullableRecordSchema + .getField("field3") + .schema() + .getTypes() + .get(0) + .getObjectProp("precision")); + assertEquals( + 2, + nullableRecordSchema.getField("field3").schema().getTypes().get(0).getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, + nullableRecordSchema.getField("field3").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.LONG, nullableRecordSchema.getField("field4").schema().getType()); + assertEquals( + LogicalTypes.timestampMillis(), + nullableRecordSchema.getField("field4").schema().getLogicalType()); + + // Assertions for nonNullableRecord + assertEquals(Schema.Type.RECORD, schema.getField("nonNullableRecord").schema().getType()); + Schema nonNullableRecordSchema = schema.getField("nonNullableRecord").schema(); + assertEquals(4, nonNullableRecordSchema.getFields().size()); + assertEquals( + Schema.Type.INT, + nonNullableRecordSchema.getField("field1").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.NULL, + nonNullableRecordSchema.getField("field1").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.DOUBLE, nonNullableRecordSchema.getField("field2").schema().getType()); + assertEquals( + Schema.Type.FIXED, + nonNullableRecordSchema.getField("field3").schema().getTypes().get(0).getType()); + assertEquals( + 16, nullableRecordSchema.getField("field3").schema().getTypes().get(0).getFixedSize()); + assertEquals( + LogicalTypes.decimal(10, 2), + nonNullableRecordSchema.getField("field3").schema().getTypes().get(0).getLogicalType()); + assertEquals( + 10, + nonNullableRecordSchema + .getField("field3") + .schema() + .getTypes() + .get(0) + .getObjectProp("precision")); + assertEquals( + 2, + nonNullableRecordSchema + .getField("field3") + .schema() + .getTypes() + .get(0) + .getObjectProp("scale")); + assertEquals( + Schema.Type.NULL, + nonNullableRecordSchema.getField("field3").schema().getTypes().get(1).getType()); + assertEquals(Schema.Type.LONG, nonNullableRecordSchema.getField("field4").schema().getType()); + assertEquals( + LogicalTypes.timestampMillis(), + nonNullableRecordSchema.getField("field4").schema().getLogicalType()); + } + + @Test + public void testConvertUnionTypes() { + List fields = + Arrays.asList( + new Field( + "sparseUnionField", + FieldType.nullable( + new ArrowType.Union( + UnionMode.Sparse, + new int[] { + ArrowType.ArrowTypeID.Int.getFlatbufID(), + ArrowType.ArrowTypeID.FloatingPoint.getFlatbufID(), + ArrowType.ArrowTypeID.Utf8.getFlatbufID() + })), + Arrays.asList( + new Field( + "intMember", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field( + "floatMember", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field("stringMember", FieldType.notNullable(new ArrowType.Utf8()), null))), + new Field( + "denseUnionField", + FieldType.nullable( + new ArrowType.Union( + UnionMode.Dense, + new int[] { + ArrowType.ArrowTypeID.Int.getFlatbufID(), + ArrowType.ArrowTypeID.FloatingPoint.getFlatbufID(), + ArrowType.ArrowTypeID.Utf8.getFlatbufID() + })), + Arrays.asList( + new Field( + "intMember", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field( + "floatMember", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field("stringMember", FieldType.notNullable(new ArrowType.Utf8()), null))), + new Field( + "nullableSparseUnionField", + FieldType.nullable( + new ArrowType.Union( + UnionMode.Sparse, + new int[] { + ArrowType.ArrowTypeID.Int.getFlatbufID(), + ArrowType.ArrowTypeID.FloatingPoint.getFlatbufID(), + ArrowType.ArrowTypeID.Utf8.getFlatbufID() + })), + Arrays.asList( + new Field( + "nullableIntMember", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "nullableFloatMember", + FieldType.nullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nullableStringMember", FieldType.nullable(new ArrowType.Utf8()), null))), + new Field( + "nullableDenseUnionField", + FieldType.nullable( + new ArrowType.Union( + UnionMode.Dense, + new int[] { + ArrowType.ArrowTypeID.Int.getFlatbufID(), + ArrowType.ArrowTypeID.FloatingPoint.getFlatbufID(), + ArrowType.ArrowTypeID.Utf8.getFlatbufID() + })), + Arrays.asList( + new Field( + "nullableIntMember", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "nullableFloatMember", + FieldType.nullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nullableStringMember", FieldType.nullable(new ArrowType.Utf8()), null)))); + + Schema schema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + + assertEquals(Schema.Type.RECORD, schema.getType()); + assertEquals(4, schema.getFields().size()); + + // Assertions for sparseUnionField + assertEquals(Schema.Type.UNION, schema.getField("sparseUnionField").schema().getType()); + assertEquals(3, schema.getField("sparseUnionField").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("sparseUnionField").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.FLOAT, + schema.getField("sparseUnionField").schema().getTypes().get(1).getType()); + assertEquals( + Schema.Type.STRING, + schema.getField("sparseUnionField").schema().getTypes().get(2).getType()); + + // Assertions for denseUnionField + assertEquals(Schema.Type.UNION, schema.getField("denseUnionField").schema().getType()); + assertEquals(3, schema.getField("denseUnionField").schema().getTypes().size()); + assertEquals( + Schema.Type.INT, schema.getField("denseUnionField").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.FLOAT, schema.getField("denseUnionField").schema().getTypes().get(1).getType()); + assertEquals( + Schema.Type.STRING, + schema.getField("denseUnionField").schema().getTypes().get(2).getType()); + + // Assertions for sparseUnionField + assertEquals(Schema.Type.UNION, schema.getField("nullableSparseUnionField").schema().getType()); + assertEquals(4, schema.getField("nullableSparseUnionField").schema().getTypes().size()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableSparseUnionField").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.INT, + schema.getField("nullableSparseUnionField").schema().getTypes().get(1).getType()); + assertEquals( + Schema.Type.FLOAT, + schema.getField("nullableSparseUnionField").schema().getTypes().get(2).getType()); + assertEquals( + Schema.Type.STRING, + schema.getField("nullableSparseUnionField").schema().getTypes().get(3).getType()); + + // Assertions for denseUnionField + assertEquals(Schema.Type.UNION, schema.getField("nullableDenseUnionField").schema().getType()); + assertEquals(4, schema.getField("nullableDenseUnionField").schema().getTypes().size()); + assertEquals( + Schema.Type.NULL, + schema.getField("nullableDenseUnionField").schema().getTypes().get(0).getType()); + assertEquals( + Schema.Type.INT, + schema.getField("nullableDenseUnionField").schema().getTypes().get(1).getType()); + assertEquals( + Schema.Type.FLOAT, + schema.getField("nullableDenseUnionField").schema().getTypes().get(2).getType()); + assertEquals( + 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/AvroLogicalTypesTest.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/AvroLogicalTypesTest.java index 173cc855b1..801456d79b 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/AvroLogicalTypesTest.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/AvroLogicalTypesTest.java @@ -173,7 +173,7 @@ public void testInvalidDecimalPrecision() throws Exception { IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> writeAndRead(schema, data)); - assertTrue(e.getMessage().contains("Precision must be in range of 1 to 38")); + assertTrue(e.getMessage().contains("Precision must be in range of 1 to 76")); } @Test 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 new file mode 100644 index 0000000000..ceaf59aa72 --- /dev/null +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java @@ -0,0 +1,1700 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZonedDateTime; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.NullVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeMilliVector; +import org.apache.arrow.vector.TimeStampMicroTZVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TimeStampMilliTZVector; +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; +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.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; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.EncoderFactory; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class RoundTripDataTest { + + @TempDir public static File TMP; + + 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, + 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, dictionaries)); + } + } + + 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(), dictionaries); + for (int row = 0; row < rowCount; row++) { + producer.produce(encoder); + } + encoder.flush(); + } + + // Generate AVRO schema + 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, roundTripDictionaries)) { + + assertEquals(root.getSchema(), roundTrip.getSchema()); + assertEquals(rowCount, roundTrip.getRowCount()); + + // Read and check values + 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)); + } + } + } + } + } + + private static void roundTripByteArrayTest( + VectorSchemaRoot root, BufferAllocator allocator, File dataFile, int rowCount) + 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()); + for (int row = 0; row < rowCount; row++) { + producer.produce(encoder); + } + encoder.flush(); + } + + // Generate AVRO schema + Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields()); + + // Read back in and compare + try (VectorSchemaRoot roundTrip = readDataFile(schema, dataFile, allocator, null)) { + + assertEquals(root.getSchema(), roundTrip.getSchema()); + assertEquals(rowCount, roundTrip.getRowCount()); + + // Read and check values + for (int row = 0; row < rowCount; row++) { + byte[] rootBytes = (byte[]) root.getVector(0).getObject(row); + byte[] roundTripBytes = (byte[]) roundTrip.getVector(0).getObject(row); + assertArrayEquals(rootBytes, roundTripBytes); + } + } + } + + // Data round trip for primitive types, nullable and non-nullable + + @Test + public void testRoundTripNullColumn() throws Exception { + + // The current read implementation expects EOF, which never happens for a single null vector + // Include a boolean vector with this test for now, so that EOF exception will be triggered + + // Field definition + FieldType nullField = new FieldType(false, new ArrowType.Null(), null); + FieldType booleanField = new FieldType(false, new ArrowType.Bool(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + NullVector nullVector = new NullVector(new Field("nullColumn", nullField, null)); + BitVector booleanVector = new BitVector(new Field("boolean", booleanField, null), allocator); + + int rowCount = 10; + + // Set up VSR + List vectors = Arrays.asList(nullVector, booleanVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set all values to null + for (int row = 0; row < rowCount; row++) { + nullVector.setNull(row); + booleanVector.set(row, 0); + } + + File dataFile = new File(TMP, "testRoundTripNullColumn.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripBooleans() throws Exception { + + // Field definition + FieldType booleanField = new FieldType(false, new ArrowType.Bool(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + BitVector booleanVector = new BitVector(new Field("boolean", booleanField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(booleanVector); + int rowCount = 10; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < rowCount; row++) { + booleanVector.set(row, row % 2 == 0 ? 1 : 0); + } + + File dataFile = new File(TMP, "testRoundTripBooleans.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableBooleans() throws Exception { + + // Field definition + FieldType booleanField = new FieldType(true, new ArrowType.Bool(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + BitVector booleanVector = new BitVector(new Field("boolean", booleanField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(booleanVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null value + booleanVector.setNull(0); + + // False value + booleanVector.set(1, 0); + + // True value + booleanVector.set(2, 1); + + File dataFile = new File(TMP, "testRoundTripNullableBooleans.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripIntegers() throws Exception { + + // Field definitions + FieldType int32Field = new FieldType(false, new ArrowType.Int(32, true), null); + FieldType int64Field = new FieldType(false, new ArrowType.Int(64, true), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + IntVector int32Vector = new IntVector(new Field("int32", int32Field, null), allocator); + BigIntVector int64Vector = new BigIntVector(new Field("int64", int64Field, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(int32Vector, int64Vector); + + int rowCount = 12; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < 10; row++) { + int32Vector.set(row, 513 * row * (row % 2 == 0 ? 1 : -1)); + int64Vector.set(row, 3791L * row * (row % 2 == 0 ? 1 : -1)); + } + + // Min values + int32Vector.set(10, Integer.MIN_VALUE); + int64Vector.set(10, Long.MIN_VALUE); + + // Max values + int32Vector.set(11, Integer.MAX_VALUE); + int64Vector.set(11, Long.MAX_VALUE); + + File dataFile = new File(TMP, "testRoundTripIntegers.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableIntegers() throws Exception { + + // Field definitions + FieldType int32Field = new FieldType(true, new ArrowType.Int(32, true), null); + FieldType int64Field = new FieldType(true, new ArrowType.Int(64, true), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + IntVector int32Vector = new IntVector(new Field("int32", int32Field, null), allocator); + BigIntVector int64Vector = new BigIntVector(new Field("int64", int64Field, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(int32Vector, int64Vector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null values + int32Vector.setNull(0); + int64Vector.setNull(0); + + // Zero values + int32Vector.set(1, 0); + int64Vector.set(1, 0); + + // Non-zero values + int32Vector.set(2, Integer.MAX_VALUE); + int64Vector.set(2, Long.MAX_VALUE); + + File dataFile = new File(TMP, "testRoundTripNullableIntegers.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripFloatingPoints() throws Exception { + + // Field definitions + FieldType float32Field = + new FieldType(false, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null); + FieldType float64Field = + new FieldType(false, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + Float4Vector float32Vector = + new Float4Vector(new Field("float32", float32Field, null), allocator); + Float8Vector float64Vector = + new Float8Vector(new Field("float64", float64Field, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(float32Vector, float64Vector); + int rowCount = 15; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + for (int row = 0; row < 10; row++) { + float32Vector.set(row, 37.6f * row * (row % 2 == 0 ? 1 : -1)); + float64Vector.set(row, 37.6d * row * (row % 2 == 0 ? 1 : -1)); + } + + float32Vector.set(10, Float.MIN_VALUE); + float64Vector.set(10, Double.MIN_VALUE); + + float32Vector.set(11, Float.MAX_VALUE); + float64Vector.set(11, Double.MAX_VALUE); + + float32Vector.set(12, Float.NaN); + float64Vector.set(12, Double.NaN); + + float32Vector.set(13, Float.POSITIVE_INFINITY); + float64Vector.set(13, Double.POSITIVE_INFINITY); + + float32Vector.set(14, Float.NEGATIVE_INFINITY); + float64Vector.set(14, Double.NEGATIVE_INFINITY); + + File dataFile = new File(TMP, "testRoundTripFloatingPoints.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableFloatingPoints() throws Exception { + + // Field definitions + FieldType float32Field = + new FieldType(true, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), null); + FieldType float64Field = + new FieldType(true, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + Float4Vector float32Vector = + new Float4Vector(new Field("float32", float32Field, null), allocator); + Float8Vector float64Vector = + new Float8Vector(new Field("float64", float64Field, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(float32Vector, float64Vector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Null values + float32Vector.setNull(0); + float64Vector.setNull(0); + + // Zero values + float32Vector.set(1, 0.0f); + float64Vector.set(1, 0.0); + + // Non-zero values + float32Vector.set(2, 1.0f); + float64Vector.set(2, 1.0); + + File dataFile = new File(TMP, "testRoundTripNullableFloatingPoints.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripStrings() throws Exception { + + // Field definition + FieldType stringField = new FieldType(false, new ArrowType.Utf8(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarCharVector stringVector = + new VarCharVector(new Field("string", stringField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(stringVector); + int rowCount = 5; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + stringVector.setSafe(0, "Hello world!".getBytes()); + stringVector.setSafe(1, "<%**\r\n\t\\abc\0$$>".getBytes()); + stringVector.setSafe(2, "你好世界".getBytes()); + stringVector.setSafe(3, "مرحبا بالعالم".getBytes()); + stringVector.setSafe(4, "(P ∧ P ⇒ Q) ⇒ Q".getBytes()); + + File dataFile = new File(TMP, "testRoundTripStrings.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableStrings() throws Exception { + + // Field definition + FieldType stringField = new FieldType(true, new ArrowType.Utf8(), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarCharVector stringVector = + new VarCharVector(new Field("string", stringField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(stringVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + stringVector.setNull(0); + stringVector.setSafe(1, "".getBytes()); + stringVector.setSafe(2, "not empty".getBytes()); + + File dataFile = new File(TMP, "testRoundTripNullableStrings.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripBinary() throws Exception { + + // Field definition + FieldType binaryField = new FieldType(false, new ArrowType.Binary(), null); + FieldType fixedField = new FieldType(false, new ArrowType.FixedSizeBinary(5), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarBinaryVector binaryVector = + new VarBinaryVector(new Field("binary", binaryField, null), allocator); + FixedSizeBinaryVector fixedVector = + new FixedSizeBinaryVector(new Field("fixed", fixedField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(binaryVector, fixedVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + binaryVector.setSafe(0, new byte[] {1, 2, 3}); + binaryVector.setSafe(1, new byte[] {4, 5, 6, 7}); + binaryVector.setSafe(2, new byte[] {8, 9}); + + fixedVector.setSafe(0, new byte[] {1, 2, 3, 4, 5}); + fixedVector.setSafe(1, new byte[] {4, 5, 6, 7, 8, 9}); + fixedVector.setSafe(2, new byte[] {8, 9, 10, 11, 12}); + + File dataFile = new File(TMP, "testRoundTripBinary.avro"); + + roundTripByteArrayTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableBinary() throws Exception { + + // Field definition + FieldType binaryField = new FieldType(true, new ArrowType.Binary(), null); + FieldType fixedField = new FieldType(true, new ArrowType.FixedSizeBinary(5), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + VarBinaryVector binaryVector = + new VarBinaryVector(new Field("binary", binaryField, null), allocator); + FixedSizeBinaryVector fixedVector = + new FixedSizeBinaryVector(new Field("fixed", fixedField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(binaryVector, fixedVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + binaryVector.setNull(0); + binaryVector.setSafe(1, new byte[] {}); + binaryVector.setSafe(2, new byte[] {10, 11, 12}); + + fixedVector.setNull(0); + fixedVector.setSafe(1, new byte[] {0, 0, 0, 0, 0}); + fixedVector.setSafe(2, new byte[] {10, 11, 12, 13, 14}); + + File dataFile = new File(TMP, "testRoundTripNullableBinary.avro"); + + roundTripByteArrayTest(root, allocator, dataFile, rowCount); + } + } + + // Data round trip for logical types, nullable and non-nullable + + @Test + public void testRoundTripDecimals() throws Exception { + + // Field definitions + FieldType decimal128Field1 = new FieldType(false, new ArrowType.Decimal(38, 10, 128), null); + FieldType decimal128Field2 = new FieldType(false, new ArrowType.Decimal(38, 5, 128), null); + FieldType decimal256Field1 = new FieldType(false, new ArrowType.Decimal(76, 20, 256), null); + FieldType decimal256Field2 = new FieldType(false, new ArrowType.Decimal(76, 10, 256), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DecimalVector decimal128Vector1 = + new DecimalVector(new Field("decimal128_1", decimal128Field1, null), allocator); + DecimalVector decimal128Vector2 = + new DecimalVector(new Field("decimal128_2", decimal128Field2, null), allocator); + Decimal256Vector decimal256Vector1 = + new Decimal256Vector(new Field("decimal256_1", decimal256Field1, null), allocator); + Decimal256Vector decimal256Vector2 = + new Decimal256Vector(new Field("decimal256_2", decimal256Field2, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList(decimal128Vector1, decimal128Vector2, decimal256Vector1, decimal256Vector2); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + decimal128Vector1.setSafe( + 0, new BigDecimal("12345.67890").setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 1, new BigDecimal("-98765.43210").setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 2, new BigDecimal("54321.09876").setScale(10, RoundingMode.UNNECESSARY)); + + decimal128Vector2.setSafe( + 0, new BigDecimal("12345.67890").setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 1, new BigDecimal("-98765.43210").setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 2, new BigDecimal("54321.09876").setScale(5, RoundingMode.UNNECESSARY)); + + decimal256Vector1.setSafe( + 0, + new BigDecimal("12345678901234567890.12345678901234567890") + .setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 1, + new BigDecimal("-98765432109876543210.98765432109876543210") + .setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 2, + new BigDecimal("54321098765432109876.54321098765432109876") + .setScale(20, RoundingMode.UNNECESSARY)); + + decimal256Vector2.setSafe( + 0, + new BigDecimal("12345678901234567890.1234567890").setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 1, + new BigDecimal("-98765432109876543210.9876543210") + .setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 2, + new BigDecimal("54321098765432109876.5432109876").setScale(10, RoundingMode.UNNECESSARY)); + + File dataFile = new File(TMP, "testRoundTripDecimals.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableDecimals() throws Exception { + + // Field definitions + FieldType decimal128Field1 = new FieldType(true, new ArrowType.Decimal(38, 10, 128), null); + FieldType decimal128Field2 = new FieldType(true, new ArrowType.Decimal(38, 5, 128), null); + FieldType decimal256Field1 = new FieldType(true, new ArrowType.Decimal(76, 20, 256), null); + FieldType decimal256Field2 = new FieldType(true, new ArrowType.Decimal(76, 10, 256), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DecimalVector decimal128Vector1 = + new DecimalVector(new Field("decimal128_1", decimal128Field1, null), allocator); + DecimalVector decimal128Vector2 = + new DecimalVector(new Field("decimal128_2", decimal128Field2, null), allocator); + Decimal256Vector decimal256Vector1 = + new Decimal256Vector(new Field("decimal256_1", decimal256Field1, null), allocator); + Decimal256Vector decimal256Vector2 = + new Decimal256Vector(new Field("decimal256_2", decimal256Field2, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList(decimal128Vector1, decimal128Vector2, decimal256Vector1, decimal256Vector2); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + decimal128Vector1.setNull(0); + decimal128Vector1.setSafe(1, BigDecimal.ZERO.setScale(10, RoundingMode.UNNECESSARY)); + decimal128Vector1.setSafe( + 2, new BigDecimal("12345.67890").setScale(10, RoundingMode.UNNECESSARY)); + + decimal128Vector2.setNull(0); + decimal128Vector2.setSafe(1, BigDecimal.ZERO.setScale(5, RoundingMode.UNNECESSARY)); + decimal128Vector2.setSafe( + 2, new BigDecimal("98765.43210").setScale(5, RoundingMode.UNNECESSARY)); + + decimal256Vector1.setNull(0); + decimal256Vector1.setSafe(1, BigDecimal.ZERO.setScale(20, RoundingMode.UNNECESSARY)); + decimal256Vector1.setSafe( + 2, + new BigDecimal("12345678901234567890.12345678901234567890") + .setScale(20, RoundingMode.UNNECESSARY)); + + decimal256Vector2.setNull(0); + decimal256Vector2.setSafe(1, BigDecimal.ZERO.setScale(10, RoundingMode.UNNECESSARY)); + decimal256Vector2.setSafe( + 2, + new BigDecimal("98765432109876543210.9876543210").setScale(10, RoundingMode.UNNECESSARY)); + + File dataFile = new File(TMP, "testRoundTripNullableDecimals.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripDates() throws Exception { + + // Field definitions + FieldType dateDayField = new FieldType(false, new ArrowType.Date(DateUnit.DAY), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DateDayVector dateDayVector = + new DateDayVector(new Field("dateDay", dateDayField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(dateDayVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + dateDayVector.setSafe(0, (int) LocalDate.now().toEpochDay()); + dateDayVector.setSafe(1, (int) LocalDate.now().toEpochDay() + 1); + dateDayVector.setSafe(2, (int) LocalDate.now().toEpochDay() + 2); + + File dataFile = new File(TMP, "testRoundTripDates.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableDates() throws Exception { + + // Field definitions + FieldType dateDayField = new FieldType(true, new ArrowType.Date(DateUnit.DAY), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + DateDayVector dateDayVector = + new DateDayVector(new Field("dateDay", dateDayField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(dateDayVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + dateDayVector.setNull(0); + dateDayVector.setSafe(1, 0); + dateDayVector.setSafe(2, (int) LocalDate.now().toEpochDay()); + + File dataFile = new File(TMP, "testRoundTripNullableDates.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripTimes() throws Exception { + + // Field definitions + FieldType timeMillisField = + new FieldType(false, new ArrowType.Time(TimeUnit.MILLISECOND, 32), null); + FieldType timeMicrosField = + new FieldType(false, new ArrowType.Time(TimeUnit.MICROSECOND, 64), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeMilliVector timeMillisVector = + new TimeMilliVector(new Field("timeMillis", timeMillisField, null), allocator); + TimeMicroVector timeMicrosVector = + new TimeMicroVector(new Field("timeMicros", timeMicrosField, null), allocator); + + // Set up VSR + List vectors = Arrays.asList(timeMillisVector, timeMicrosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timeMillisVector.setSafe( + 0, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000)); + timeMillisVector.setSafe( + 1, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000) - 1000); + timeMillisVector.setSafe( + 2, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000) - 2000); + + timeMicrosVector.setSafe(0, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000); + timeMicrosVector.setSafe(1, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000 - 1000000); + timeMicrosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000 - 2000000); + + File dataFile = new File(TMP, "testRoundTripTimes.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableTimes() throws Exception { + + // Field definitions + FieldType timeMillisField = + new FieldType(true, new ArrowType.Time(TimeUnit.MILLISECOND, 32), null); + FieldType timeMicrosField = + new FieldType(true, new ArrowType.Time(TimeUnit.MICROSECOND, 64), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeMilliVector timeMillisVector = + new TimeMilliVector(new Field("timeMillis", timeMillisField, null), allocator); + TimeMicroVector timeMicrosVector = + new TimeMicroVector(new Field("timeMicros", timeMicrosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = Arrays.asList(timeMillisVector, timeMicrosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timeMillisVector.setNull(0); + timeMillisVector.setSafe(1, 0); + timeMillisVector.setSafe( + 2, (int) (ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000000)); + + timeMicrosVector.setNull(0); + timeMicrosVector.setSafe(1, 0); + timeMicrosVector.setSafe(2, ZonedDateTime.now().toLocalTime().toNanoOfDay() / 1000); + + File dataFile = new File(TMP, "testRoundTripNullableTimes.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripZoneAwareTimestamps() throws Exception { + + // Field definitions + FieldType timestampMillisField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"), null); + FieldType timestampMicrosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), null); + FieldType timestampNanosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampMilliTZVector timestampMillisVector = + new TimeStampMilliTZVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroTZVector timestampMicrosVector = + new TimeStampMicroTZVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoTZVector timestampNanosVector = + new TimeStampNanoTZVector( + new Field("timestampNanos", timestampNanosField, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList(timestampMillisVector, timestampMicrosVector, timestampNanosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampMillisVector.setSafe(0, (int) Instant.now().toEpochMilli()); + timestampMillisVector.setSafe(1, (int) Instant.now().toEpochMilli() - 1000); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli() - 2000); + + timestampMicrosVector.setSafe(0, Instant.now().toEpochMilli() * 1000); + timestampMicrosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000); + timestampMicrosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000); + + timestampNanosVector.setSafe(0, Instant.now().toEpochMilli() * 1000000); + timestampNanosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000000); + timestampNanosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000000); + + File dataFile = new File(TMP, "testRoundTripZoneAwareTimestamps.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableZoneAwareTimestamps() throws Exception { + + // Field definitions + FieldType timestampMillisField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"), null); + FieldType timestampMicrosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"), null); + FieldType timestampNanosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampMilliTZVector timestampMillisVector = + new TimeStampMilliTZVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroTZVector timestampMicrosVector = + new TimeStampMicroTZVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoTZVector timestampNanosVector = + new TimeStampNanoTZVector( + new Field("timestampNanos", timestampNanosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList(timestampMillisVector, timestampMicrosVector, timestampNanosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampMillisVector.setNull(0); + timestampMillisVector.setSafe(1, 0); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli()); + + timestampMicrosVector.setNull(0); + timestampMicrosVector.setSafe(1, 0); + timestampMicrosVector.setSafe(2, Instant.now().toEpochMilli() * 1000); + + timestampNanosVector.setNull(0); + timestampNanosVector.setSafe(1, 0); + timestampNanosVector.setSafe(2, Instant.now().toEpochMilli() * 1000000); + + File dataFile = new File(TMP, "testRoundTripNullableZoneAwareTimestamps.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripLocalTimestamps() throws Exception { + + // Field definitions + FieldType timestampMillisField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null); + FieldType timestampMicrosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), null); + FieldType timestampNanosField = + new FieldType(false, new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampMilliVector timestampMillisVector = + new TimeStampMilliVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroVector timestampMicrosVector = + new TimeStampMicroVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoVector timestampNanosVector = + new TimeStampNanoVector(new Field("timestampNanos", timestampNanosField, null), allocator); + + // Set up VSR + List vectors = + Arrays.asList(timestampMillisVector, timestampMicrosVector, timestampNanosVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampMillisVector.setSafe(0, (int) Instant.now().toEpochMilli()); + timestampMillisVector.setSafe(1, (int) Instant.now().toEpochMilli() - 1000); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli() - 2000); + + timestampMicrosVector.setSafe(0, Instant.now().toEpochMilli() * 1000); + timestampMicrosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000); + timestampMicrosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000); + + timestampNanosVector.setSafe(0, Instant.now().toEpochMilli() * 1000000); + timestampNanosVector.setSafe(1, (Instant.now().toEpochMilli() - 1000) * 1000000); + timestampNanosVector.setSafe(2, (Instant.now().toEpochMilli() - 2000) * 1000000); + + File dataFile = new File(TMP, "testRoundTripLocalTimestamps.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableLocalTimestamps() throws Exception { + + // Field definitions + FieldType timestampMillisField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MILLISECOND, null), null); + FieldType timestampMicrosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), null); + FieldType timestampNanosField = + new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND, null), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + TimeStampMilliVector timestampMillisVector = + new TimeStampMilliVector( + new Field("timestampMillis", timestampMillisField, null), allocator); + TimeStampMicroVector timestampMicrosVector = + new TimeStampMicroVector( + new Field("timestampMicros", timestampMicrosField, null), allocator); + TimeStampNanoVector timestampNanosVector = + new TimeStampNanoVector(new Field("timestampNanos", timestampNanosField, null), allocator); + + int rowCount = 3; + + // Set up VSR + List vectors = + Arrays.asList(timestampMillisVector, timestampMicrosVector, timestampNanosVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + timestampMillisVector.setNull(0); + timestampMillisVector.setSafe(1, 0); + timestampMillisVector.setSafe(2, (int) Instant.now().toEpochMilli()); + + timestampMicrosVector.setNull(0); + timestampMicrosVector.setSafe(1, 0); + timestampMicrosVector.setSafe(2, Instant.now().toEpochMilli() * 1000); + + timestampNanosVector.setNull(0); + timestampNanosVector.setSafe(1, 0); + timestampNanosVector.setSafe(2, Instant.now().toEpochMilli() * 1000000); + + File dataFile = new File(TMP, "testRoundTripNullableLocalTimestamps.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + // Data round trip for containers of primitive and logical types, nullable and non-nullable + + @Test + public void testRoundTripLists() throws Exception { + + // Field definitions + FieldType intListField = new FieldType(false, new ArrowType.List(), null); + FieldType stringListField = new FieldType(false, new ArrowType.List(), null); + FieldType dateListField = new FieldType(false, new ArrowType.List(), null); + + Field intField = new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("item", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("item", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + ListVector intListVector = new ListVector("intList", allocator, intListField, null); + ListVector stringListVector = new ListVector("stringList", allocator, stringListField, null); + ListVector dateListVector = new ListVector("dateList", allocator, dateListField, null); + + intListVector.initializeChildrenFromFields(Arrays.asList(intField)); + stringListVector.initializeChildrenFromFields(Arrays.asList(stringField)); + dateListVector.initializeChildrenFromFields(Arrays.asList(dateField)); + + // Set up VSR + List vectors = Arrays.asList(intListVector, stringListVector, dateListVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + FieldWriter intListWriter = intListVector.getWriter(); + FieldWriter stringListWriter = stringListVector.getWriter(); + FieldWriter dateListWriter = dateListVector.getWriter(); + + // Set test data for intList + for (int i = 0; i < rowCount; i++) { + intListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + intListWriter.writeInt(j); + } + intListWriter.endList(); + } + + // Set test data for stringList + for (int i = 0; i < rowCount; i++) { + stringListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + stringListWriter.writeVarChar("string" + j); + } + stringListWriter.endList(); + } + + // Set test data for dateList + for (int i = 0; i < rowCount; i++) { + dateListWriter.startList(); + for (int j = 0; j < 5 - i; j++) { + dateListWriter.writeDateDay((int) LocalDate.now().plusDays(j).toEpochDay()); + } + dateListWriter.endList(); + } + + // Update count for the vectors + intListVector.setValueCount(rowCount); + stringListVector.setValueCount(rowCount); + dateListVector.setValueCount(rowCount); + + File dataFile = new File(TMP, "testRoundTripLists.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableLists() throws Exception { + + // Field definitions + FieldType nullListType = new FieldType(true, new ArrowType.List(), null); + FieldType nonNullListType = new FieldType(false, new ArrowType.List(), null); + + Field nullFieldType = new Field("item", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field nonNullFieldType = + new Field("item", FieldType.notNullable(new ArrowType.Int(32, true)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + ListVector nullEntriesVector = + new ListVector("nullEntriesVector", allocator, nonNullListType, null); + ListVector nullListVector = new ListVector("nullListVector", allocator, nullListType, null); + ListVector nullBothVector = new ListVector("nullBothVector", allocator, nullListType, null); + + nullEntriesVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + nullListVector.initializeChildrenFromFields(Arrays.asList(nonNullFieldType)); + nullBothVector.initializeChildrenFromFields(Arrays.asList(nullFieldType)); + + // Set up VSR + List vectors = Arrays.asList(nullEntriesVector, nullListVector, nullBothVector); + int rowCount = 4; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for nullEntriesVector + FieldWriter nullEntriesWriter = nullEntriesVector.getWriter(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.integer().writeNull(); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.integer().writeInt(0); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.endList(); + + // Set test data for nullListVector + FieldWriter nullListWriter = nullListVector.getWriter(); + nullListWriter.writeNull(); + nullListWriter.setPosition(1); // writeNull() does not inc. idx() on list vector + nullListWriter.startList(); + nullListWriter.integer().writeInt(0); + nullListWriter.integer().writeInt(0); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.integer().writeInt(789); + nullEntriesWriter.endList(); + + // Set test data for nullBothVector + FieldWriter nullBothWriter = nullBothVector.getWriter(); + nullBothWriter.writeNull(); + nullBothWriter.setPosition(1); + nullBothWriter.startList(); + nullBothWriter.integer().writeNull(); + nullBothWriter.integer().writeNull(); + nullBothWriter.endList(); + nullListWriter.startList(); + nullListWriter.integer().writeInt(0); + nullListWriter.integer().writeInt(0); + nullListWriter.endList(); + nullEntriesWriter.startList(); + nullEntriesWriter.integer().writeInt(123); + nullEntriesWriter.integer().writeInt(456); + nullEntriesWriter.endList(); + + // Update count for the vectors + nullListVector.setValueCount(4); + nullEntriesVector.setValueCount(4); + nullBothVector.setValueCount(4); + + File dataFile = new File(TMP, "testRoundTripNullableLists.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripMap() throws Exception { + + // Field definitions + FieldType intMapField = new FieldType(false, new ArrowType.Map(false), null); + FieldType stringMapField = new FieldType(false, new ArrowType.Map(false), null); + FieldType dateMapField = new FieldType(false, new ArrowType.Map(false), null); + + Field keyField = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field intField = new Field("value", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("value", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("value", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + Field intEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, intField)); + Field stringEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, stringField)); + Field dateEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, dateField)); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + MapVector intMapVector = new MapVector("intMap", allocator, intMapField, null); + MapVector stringMapVector = new MapVector("stringMap", allocator, stringMapField, null); + MapVector dateMapVector = new MapVector("dateMap", allocator, dateMapField, null); + + intMapVector.initializeChildrenFromFields(Arrays.asList(intEntryField)); + stringMapVector.initializeChildrenFromFields(Arrays.asList(stringEntryField)); + dateMapVector.initializeChildrenFromFields(Arrays.asList(dateEntryField)); + + // Set up VSR + List vectors = Arrays.asList(intMapVector, stringMapVector, dateMapVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Total number of entries that will be writen to each vector + int entryCount = 5 + 4 + 3; + + // Set test data for intList + BaseWriter.MapWriter writer = intMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + writer.startMap(); + for (int j = 0; j < 5 - i; j++) { + writer.startEntry(); + writer.key().varChar().writeVarChar("key" + j); + writer.value().integer().writeInt(j); + writer.endEntry(); + } + writer.endMap(); + } + + // Update count for data vector (map writer does not do this) + intMapVector.getDataVector().setValueCount(entryCount); + + // Set test data for stringList + BaseWriter.MapWriter stringWriter = stringMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + stringWriter.startMap(); + for (int j = 0; j < 5 - i; j++) { + stringWriter.startEntry(); + stringWriter.key().varChar().writeVarChar("key" + j); + stringWriter.value().varChar().writeVarChar("string" + j); + stringWriter.endEntry(); + } + stringWriter.endMap(); + } + + // Update count for the vectors + intMapVector.setValueCount(rowCount); + stringMapVector.setValueCount(rowCount); + dateMapVector.setValueCount(rowCount); + + // Update count for data vector (map writer does not do this) + stringMapVector.getDataVector().setValueCount(entryCount); + + // Set test data for dateList + BaseWriter.MapWriter dateWriter = dateMapVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + dateWriter.startMap(); + for (int j = 0; j < 5 - i; j++) { + dateWriter.startEntry(); + dateWriter.key().varChar().writeVarChar("key" + j); + dateWriter.value().dateDay().writeDateDay((int) LocalDate.now().plusDays(j).toEpochDay()); + dateWriter.endEntry(); + } + dateWriter.endMap(); + } + + // Update count for data vector (map writer does not do this) + dateMapVector.getDataVector().setValueCount(entryCount); + + File dataFile = new File(TMP, "testRoundTripMap.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableMap() throws Exception { + + // Field definitions + FieldType nullMapType = new FieldType(true, new ArrowType.Map(false), null); + FieldType nonNullMapType = new FieldType(false, new ArrowType.Map(false), null); + + Field keyField = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field nullFieldType = new Field("value", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field nonNullFieldType = + new Field("value", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field nullEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, nullFieldType)); + Field nonNullEntryField = + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList(keyField, nonNullFieldType)); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + MapVector nullEntriesVector = + new MapVector("nullEntriesVector", allocator, nonNullMapType, null); + MapVector nullMapVector = new MapVector("nullMapVector", allocator, nullMapType, null); + MapVector nullBothVector = new MapVector("nullBothVector", allocator, nullMapType, null); + + nullEntriesVector.initializeChildrenFromFields(Arrays.asList(nullEntryField)); + nullMapVector.initializeChildrenFromFields(Arrays.asList(nonNullEntryField)); + nullBothVector.initializeChildrenFromFields(Arrays.asList(nullEntryField)); + + // Set up VSR + List vectors = Arrays.asList(nullEntriesVector, nullMapVector, nullBothVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for intList + BaseWriter.MapWriter writer = nullEntriesVector.getWriter(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key0"); + writer.value().integer().writeNull(); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key1"); + writer.value().integer().writeNull(); + writer.endEntry(); + writer.endMap(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key2"); + writer.value().integer().writeInt(0); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key3"); + writer.value().integer().writeInt(0); + writer.endEntry(); + writer.endMap(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + writer.endMap(); + + // Set test data for stringList + BaseWriter.MapWriter nullMapWriter = nullMapVector.getWriter(); + nullMapWriter.writeNull(); + nullMapWriter.setPosition(1); // writeNull() does not inc. idx() on map (list) vector + nullMapWriter.startMap(); + nullMapWriter.startEntry(); + nullMapWriter.key().varChar().writeVarChar("key2"); + nullMapWriter.value().integer().writeInt(0); + nullMapWriter.endEntry(); + writer.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key3"); + writer.value().integer().writeInt(0); + writer.endEntry(); + nullMapWriter.endMap(); + nullMapWriter.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + nullMapWriter.endMap(); + + // Set test data for dateList + BaseWriter.MapWriter nullBothWriter = nullBothVector.getWriter(); + nullBothWriter.writeNull(); + nullBothWriter.setPosition(1); + nullBothWriter.startMap(); + nullBothWriter.startEntry(); + nullBothWriter.key().varChar().writeVarChar("key2"); + nullBothWriter.value().integer().writeNull(); + nullBothWriter.endEntry(); + nullBothWriter.startEntry(); + nullBothWriter.key().varChar().writeVarChar("key3"); + nullBothWriter.value().integer().writeNull(); + nullBothWriter.endEntry(); + nullBothWriter.endMap(); + nullBothWriter.startMap(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key4"); + writer.value().integer().writeInt(123); + writer.endEntry(); + writer.startEntry(); + writer.key().varChar().writeVarChar("key5"); + writer.value().integer().writeInt(456); + writer.endEntry(); + nullBothWriter.endMap(); + + // Update count for the vectors + nullEntriesVector.setValueCount(3); + nullMapVector.setValueCount(3); + nullBothVector.setValueCount(3); + + File dataFile = new File(TMP, "testRoundTripNullableMap.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripStruct() throws Exception { + + // Field definitions + FieldType structFieldType = new FieldType(false, new ArrowType.Struct(), null); + Field intField = + new Field("intField", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field stringField = new Field("stringField", FieldType.notNullable(new ArrowType.Utf8()), null); + Field dateField = + new Field("dateField", FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), null); + + // Create empty vector + BufferAllocator allocator = new RootAllocator(); + StructVector structVector = new StructVector("struct", allocator, structFieldType, null); + structVector.initializeChildrenFromFields(Arrays.asList(intField, stringField, dateField)); + + // Set up VSR + List vectors = Arrays.asList(structVector); + int rowCount = 3; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data + BaseWriter.StructWriter structWriter = structVector.getWriter(); + + for (int i = 0; i < rowCount; i++) { + structWriter.start(); + structWriter.integer("intField").writeInt(i); + structWriter.varChar("stringField").writeVarChar("string" + i); + structWriter.dateDay("dateField").writeDateDay((int) LocalDate.now().toEpochDay() + i); + structWriter.end(); + } + + File dataFile = new File(TMP, "testRoundTripStruct.avro"); + + roundTripTest(root, allocator, dataFile, rowCount); + } + } + + @Test + public void testRoundTripNullableStructs() throws Exception { + + // Field definitions + FieldType structFieldType = new FieldType(false, new ArrowType.Struct(), null); + FieldType nullableStructFieldType = new FieldType(true, new ArrowType.Struct(), null); + Field intField = + new Field("intField", FieldType.notNullable(new ArrowType.Int(32, true)), null); + Field nullableIntField = + new Field("nullableIntField", FieldType.nullable(new ArrowType.Int(32, true)), null); + + // Create empty vectors + BufferAllocator allocator = new RootAllocator(); + StructVector structVector = new StructVector("struct", allocator, structFieldType, null); + StructVector nullableStructVector = + new StructVector("nullableStruct", allocator, nullableStructFieldType, null); + structVector.initializeChildrenFromFields(Arrays.asList(intField, nullableIntField)); + nullableStructVector.initializeChildrenFromFields(Arrays.asList(intField, nullableIntField)); + + // Set up VSR + List vectors = Arrays.asList(structVector, nullableStructVector); + int rowCount = 4; + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + + root.setRowCount(rowCount); + root.allocateNew(); + + // Set test data for structVector + BaseWriter.StructWriter structWriter = structVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + structWriter.setPosition(i); + structWriter.start(); + structWriter.integer("intField").writeInt(i); + if (i % 2 == 0) { + structWriter.integer("nullableIntField").writeInt(i * 10); + } else { + structWriter.integer("nullableIntField").writeNull(); + } + structWriter.end(); + } + + // Set test data for nullableStructVector + BaseWriter.StructWriter nullableStructWriter = nullableStructVector.getWriter(); + for (int i = 0; i < rowCount; i++) { + nullableStructWriter.setPosition(i); + if (i >= 2) { + nullableStructWriter.start(); + nullableStructWriter.integer("intField").writeInt(i); + if (i % 2 == 0) { + nullableStructWriter.integer("nullableIntField").writeInt(i * 10); + } else { + nullableStructWriter.integer("nullableIntField").writeNull(); + } + nullableStructWriter.end(); + } else { + nullableStructWriter.writeNull(); + } + } + + // Update count for the vector + structVector.setValueCount(rowCount); + nullableStructVector.setValueCount(rowCount); + + File dataFile = new File(TMP, "testRoundTripNullableStructs.avro"); + + 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 new file mode 100644 index 0000000000..37c0b4d9fe --- /dev/null +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java @@ -0,0 +1,500 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +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); + } + + private void doRoundTripTest(List fields, DictionaryProvider dictionaries) { + + 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, 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 + + @Test + public void testRoundTripNullType() { + + List fields = + Arrays.asList(new Field("nullType", FieldType.notNullable(new ArrowType.Null()), null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripBooleanType() { + + List fields = + Arrays.asList( + new Field("nullableBool", FieldType.nullable(new ArrowType.Bool()), null), + new Field("nonNullableBool", FieldType.notNullable(new ArrowType.Bool()), null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripIntegerTypes() { + + AvroToArrowConfig config = new AvroToArrowConfig(null, 1, null, Collections.emptySet(), false); + + // Only round trip types with direct equivalent in Avro + + List fields = + Arrays.asList( + new Field("nullableInt32", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("nonNullableInt32", FieldType.notNullable(new ArrowType.Int(32, true)), null), + new Field("nullableInt64", FieldType.nullable(new ArrowType.Int(64, true)), null), + new Field( + "nonNullableInt64", FieldType.notNullable(new ArrowType.Int(64, true)), null)); + + Schema avroSchema = ArrowToAvroUtils.createAvroSchema(fields, "TestRecord"); + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + AvroToArrowUtils.createArrowSchema(avroSchema, config); + + // Exact match on fields after round trip + assertEquals(fields, arrowSchema.getFields()); + } + + @Test + public void testRoundTripFloatingPointTypes() { + + // Only round trip types with direct equivalent in Avro + + List fields = + Arrays.asList( + new Field( + "nullableFloat32", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nonNullableFloat32", + FieldType.notNullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + null), + new Field( + "nullableFloat64", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "nonNullableFloat64", + FieldType.notNullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripStringTypes() { + + List fields = + Arrays.asList( + new Field("nullableUtf8", FieldType.nullable(new ArrowType.Utf8()), null), + new Field("nonNullableUtf8", FieldType.notNullable(new ArrowType.Utf8()), null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripBinaryTypes() { + + List fields = + Arrays.asList( + new Field("nullableBinary", FieldType.nullable(new ArrowType.Binary()), null), + new Field("nonNullableBinary", FieldType.notNullable(new ArrowType.Binary()), null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripFixedSizeBinaryTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableFixedSizeBinary", + FieldType.nullable(new ArrowType.FixedSizeBinary(10)), + null), + new Field( + "nonNullableFixedSizeBinary", + FieldType.notNullable(new ArrowType.FixedSizeBinary(10)), + null)); + + doRoundTripTest(fields); + } + + // Schema round trip for logical types, nullable and non-nullable + + @Test + public void testRoundTripDecimalTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableDecimal128", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "nonNullableDecimal1281", + FieldType.notNullable(new ArrowType.Decimal(10, 2, 128)), + null), + new Field( + "nonNullableDecimal1282", + FieldType.notNullable(new ArrowType.Decimal(15, 5, 128)), + null), + new Field( + "nonNullableDecimal1283", + FieldType.notNullable(new ArrowType.Decimal(20, 10, 128)), + null), + new Field( + "nullableDecimal256", FieldType.nullable(new ArrowType.Decimal(55, 15, 256)), null), + new Field( + "nonNullableDecimal2561", + FieldType.notNullable(new ArrowType.Decimal(55, 25, 256)), + null), + new Field( + "nonNullableDecimal2562", + FieldType.notNullable(new ArrowType.Decimal(25, 8, 256)), + null), + new Field( + "nonNullableDecimal2563", + FieldType.notNullable(new ArrowType.Decimal(60, 50, 256)), + null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripDateTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableDateDay", FieldType.nullable(new ArrowType.Date(DateUnit.DAY)), null), + new Field( + "nonNullableDateDay", + FieldType.notNullable(new ArrowType.Date(DateUnit.DAY)), + null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripTimeTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableTimeMillis", + FieldType.nullable(new ArrowType.Time(TimeUnit.MILLISECOND, 32)), + null), + new Field( + "nonNullableTimeMillis", + FieldType.notNullable(new ArrowType.Time(TimeUnit.MILLISECOND, 32)), + null), + new Field( + "nullableTimeMicros", + FieldType.nullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), + null), + new Field( + "nonNullableTimeMicros", + FieldType.notNullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), + null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripZoneAwareTimestampTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableTimestampMillisTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null), + new Field( + "nonNullableTimestampMillisTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null), + new Field( + "nullableTimestampMicrosTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null), + new Field( + "nonNullableTimestampMicrosTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null), + new Field( + "nullableTimestampNanosTz", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), + null), + new Field( + "nonNullableTimestampNanosTz", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC")), + null)); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripLocalTimestampTypes() { + + List fields = + Arrays.asList( + new Field( + "nullableTimestampMillis", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + null), + new Field( + "nonNullableTimestampMillis", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null)), + null), + new Field( + "nullableTimestampMicros", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), + null), + new Field( + "nonNullableTimestampMicros", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)), + null), + new Field( + "nullableTimestampNanos", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), + null), + new Field( + "nonNullableTimestampNanos", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.NANOSECOND, null)), + null)); + + doRoundTripTest(fields); + } + + // Schema round trip for complex types, where the contents are primitive and logical types + + @Test + public void testRoundTripListType() { + + List fields = + Arrays.asList( + new Field( + "nullableIntList", + FieldType.nullable(new ArrowType.List()), + Arrays.asList( + new Field("$data$", FieldType.nullable(new ArrowType.Int(32, true)), null))), + new Field( + "nullableDoubleList", + FieldType.nullable(new ArrowType.List()), + Arrays.asList( + new Field( + "$data$", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null))), + new Field( + "nonNullableDecimalList", + FieldType.notNullable(new ArrowType.List()), + Arrays.asList( + new Field( + "$data$", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null))), + new Field( + "nonNullableTimestampList", + FieldType.notNullable(new ArrowType.List()), + Arrays.asList( + new Field( + "$data$", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripMapType() { + + List fields = + Arrays.asList( + new Field( + "nullableMapWithNullableInt", + FieldType.nullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", FieldType.nullable(new ArrowType.Int(32, true)), null))))), + new Field( + "nullableMapWithNonNullableDouble", + FieldType.nullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null))))), + new Field( + "nonNullableMapWithNullableDecimal", + FieldType.notNullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), + null))))), + new Field( + "nonNullableMapWithNonNullableTimestamp", + FieldType.notNullable(new ArrowType.Map(false)), + Arrays.asList( + new Field( + "entries", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field( + "value", + FieldType.notNullable( + new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))))); + + doRoundTripTest(fields); + } + + @Test + public void testRoundTripStructType() { + + List fields = + Arrays.asList( + new Field( + "nullableRecord", + FieldType.nullable(new ArrowType.Struct()), + Arrays.asList( + new Field("field1", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "field2", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "field3", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "field4", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null))), + new Field( + "nonNullableRecord", + FieldType.notNullable(new ArrowType.Struct()), + Arrays.asList( + new Field("field1", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field( + "field2", + FieldType.notNullable( + new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null), + new Field( + "field3", FieldType.nullable(new ArrowType.Decimal(10, 2, 128)), null), + new Field( + "field4", + FieldType.notNullable(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC")), + null)))); + + 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/avro/src/test/java/org/apache/arrow/adapter/avro/TestWriteReadAvroRecord.java b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/TestWriteReadAvroRecord.java index c318214f5c..76e58a75ae 100644 --- a/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/TestWriteReadAvroRecord.java +++ b/adapter/avro/src/test/java/org/apache/arrow/adapter/avro/TestWriteReadAvroRecord.java @@ -19,8 +19,22 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; +import org.apache.arrow.adapter.avro.consumers.CompositeAvroConsumer; +import org.apache.arrow.adapter.avro.producers.CompositeAvroProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +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.avro.Schema; import org.apache.avro.file.DataFileReader; import org.apache.avro.file.DataFileWriter; @@ -28,10 +42,16 @@ import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericDatumWriter; import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.BinaryEncoder; import org.apache.avro.io.DatumReader; import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.EncoderFactory; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; public class TestWriteReadAvroRecord { @@ -82,4 +102,85 @@ public void testWriteAndRead() throws Exception { assertEquals(7, deUser2.get("favorite_number")); assertEquals("red", deUser2.get("favorite_color").toString()); } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testWriteAndReadVSR(boolean useSchemaFile) throws Exception { + + BufferAllocator allocator = new RootAllocator(); + FieldType stringNotNull = new FieldType(false, ArrowType.Utf8.INSTANCE, null); + FieldType stringNull = new FieldType(true, ArrowType.Utf8.INSTANCE, null); + FieldType intN32Null = new FieldType(true, new ArrowType.Int(32, true), null); + + List fields = new ArrayList<>(); + fields.add(new Field("name", stringNotNull, null)); + fields.add(new Field("favorite_number", intN32Null, null)); + fields.add(new Field("favorite_color", stringNull, null)); + + VarCharVector nameVector = new VarCharVector(fields.get(0), allocator); + nameVector.allocateNew(2); + nameVector.set(0, "Alyssa".getBytes(StandardCharsets.UTF_8)); + nameVector.set(1, "Ben".getBytes(StandardCharsets.UTF_8)); + + IntVector favNumberVector = new IntVector(fields.get(1), allocator); + favNumberVector.allocateNew(2); + favNumberVector.set(0, 256); + favNumberVector.set(1, 7); + + VarCharVector favColorVector = new VarCharVector(fields.get(2), allocator); + favColorVector.allocateNew(2); + favColorVector.setNull(0); + favColorVector.set(1, "red".getBytes(StandardCharsets.UTF_8)); + + List vectors = new ArrayList<>(); + vectors.add(nameVector); + vectors.add(favNumberVector); + vectors.add(favColorVector); + + Schema schema = + useSchemaFile + ? AvroTestBase.getSchema("test.avsc") + : ArrowToAvroUtils.createAvroSchema(fields); + + File dataFile = new File(TMP, "test_vsr.avro"); + AvroToArrowConfig config = new AvroToArrowConfigBuilder(allocator).build(); + + try (FileOutputStream fos = new FileOutputStream(dataFile)) { + + BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null); + CompositeAvroProducer producer = ArrowToAvroUtils.createCompositeProducer(vectors); + + producer.produce(encoder); + producer.produce(encoder); + + encoder.flush(); + } + + List roundTripFields = new ArrayList<>(); + List roundTripVectors = new ArrayList<>(); + + try (FileInputStream fis = new FileInputStream(dataFile)) { + + BinaryDecoder decoder = new DecoderFactory().directBinaryDecoder(fis, null); + CompositeAvroConsumer consumer = AvroToArrowUtils.createCompositeConsumer(schema, config); + + consumer.getConsumers().forEach(c -> roundTripFields.add(c.getVector().getField())); + consumer.getConsumers().forEach(c -> roundTripVectors.add(c.getVector())); + consumer.consume(decoder); + consumer.consume(decoder); + } + + VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 2); + VectorSchemaRoot roundTripRoot = new VectorSchemaRoot(roundTripFields, roundTripVectors, 2); + + assertEquals(root.getRowCount(), roundTripRoot.getRowCount()); + + for (int row = 0; row < 2; row++) { + for (int col = 0; col < 3; col++) { + FieldVector vector = root.getVector(col); + FieldVector roundTripVector = roundTripRoot.getVector(col); + assertEquals(vector.getObject(row), roundTripVector.getObject(row)); + } + } + } } diff --git a/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid1.avsc b/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid1.avsc index 18d7d63fc7..c1867811c7 100644 --- a/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid1.avsc +++ b/adapter/avro/src/test/resources/schema/logical/test_decimal_invalid1.avsc @@ -20,6 +20,6 @@ "name": "test", "type": "bytes", "logicalType" : "decimal", - "precision": 39, + "precision": 77, "scale": 2 } diff --git a/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_micros.avsc b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_micros.avsc new file mode 100644 index 0000000000..db456e8a84 --- /dev/null +++ b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_micros.avsc @@ -0,0 +1,23 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +{ + "namespace": "org.apache.arrow.avro", + "name": "test", + "type": "long", + "logicalType" : "local-timestamp-micros" +} diff --git a/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_millis.avsc b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_millis.avsc new file mode 100644 index 0000000000..6a3cf9bccb --- /dev/null +++ b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_millis.avsc @@ -0,0 +1,23 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +{ + "namespace": "org.apache.arrow.avro", + "name": "test", + "type": "long", + "logicalType" : "local-timestamp-millis" +} diff --git a/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_nanos.avsc b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_nanos.avsc new file mode 100644 index 0000000000..96ca8bbfa4 --- /dev/null +++ b/adapter/avro/src/test/resources/schema/logical/test_local_timestamp_nanos.avsc @@ -0,0 +1,23 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +{ + "namespace": "org.apache.arrow.avro", + "name": "test", + "type": "long", + "logicalType" : "local-timestamp-nanos" +} diff --git a/adapter/avro/src/test/resources/schema/logical/test_timestamp_nanos.avsc b/adapter/avro/src/test/resources/schema/logical/test_timestamp_nanos.avsc new file mode 100644 index 0000000000..9e05eab408 --- /dev/null +++ b/adapter/avro/src/test/resources/schema/logical/test_timestamp_nanos.avsc @@ -0,0 +1,23 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +{ + "namespace": "org.apache.arrow.avro", + "name": "test", + "type": "long", + "logicalType" : "timestamp-nanos" +} 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/JdbcParameterBinder.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinder.java index fd4721bcd9..d41dcb91bd 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinder.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinder.java @@ -116,7 +116,7 @@ public static class Builder { /** Bind each column to the corresponding parameter in order. */ public Builder bindAll() { for (int i = 0; i < root.getFieldVectors().size(); i++) { - bind(/*parameterIndex=*/ i + 1, /*columnIndex=*/ i); + bind(/* parameterIndex= */ i + 1, /* columnIndex= */ i); } return this; } diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java index 1bfcfc8fe0..2992bdb251 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowConfig.java @@ -61,6 +61,7 @@ public final class JdbcToArrowConfig { private final Map schemaMetadata; private final Map> columnMetadataByColumnIndex; private final RoundingMode bigDecimalRoundingMode; + /** * The maximum rowCount to read each time when partially convert data. Default value is 1024 and * -1 means disable partial read. default is -1 which means disable partial read. Note that this diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java index aecb734a8b..1edb6261d6 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java @@ -145,7 +145,7 @@ public static Schema jdbcToArrowSchema( final int scale = parameterMetaData.getScale(parameterCounter); final ArrowType arrowType = getArrowTypeFromJdbcType(new JdbcFieldInfo(jdbcDataType, precision, scale), calendar); - final FieldType fieldType = new FieldType(arrowIsNullable, arrowType, /*dictionary=*/ null); + final FieldType fieldType = new FieldType(arrowIsNullable, arrowType, /* dictionary= */ null); parameterFields.add(new Field(null, fieldType, null)); } diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumer.java index edbc6360df..73ec04b8a0 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumer.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumer.java @@ -51,13 +51,15 @@ public BinaryConsumer(VarBinaryVector vector, int index) { /** consume a InputStream. */ public void consume(InputStream is) throws IOException { + while (currentIndex >= vector.getValueCapacity()) { + vector.reallocValidityAndOffsetBuffers(); + } + + final int startOffset = vector.getStartOffset(currentIndex); + final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); + int dataLength = 0; + if (is != null) { - while (currentIndex >= vector.getValueCapacity()) { - vector.reallocValidityAndOffsetBuffers(); - } - final int startOffset = vector.getStartOffset(currentIndex); - final ArrowBuf offsetBuffer = vector.getOffsetBuffer(); - int dataLength = 0; int read; while ((read = is.read(reuseBytes)) != -1) { while (vector.getDataBuffer().capacity() < (startOffset + dataLength + read)) { @@ -66,11 +68,12 @@ public void consume(InputStream is) throws IOException { vector.getDataBuffer().setBytes(startOffset + dataLength, reuseBytes, 0, read); dataLength += read; } - offsetBuffer.setInt( - (currentIndex + 1) * ((long) VarBinaryVector.OFFSET_WIDTH), startOffset + dataLength); + BitVectorHelper.setBit(vector.getValidityBuffer(), currentIndex); - vector.setLastSet(currentIndex); } + offsetBuffer.setInt( + (currentIndex + 1) * ((long) VarBinaryVector.OFFSET_WIDTH), startOffset + dataLength); + vector.setLastSet(currentIndex); } public void moveWriterPosition() { @@ -95,9 +98,7 @@ public NullableBinaryConsumer(VarBinaryVector vector, int index) { @Override public void consume(ResultSet resultSet) throws SQLException, IOException { InputStream is = resultSet.getBinaryStream(columnIndexInResultSet); - if (!resultSet.wasNull()) { - consume(is); - } + consume(is); moveWriterPosition(); } } 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/JdbcParameterBinderTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java index a05130f18e..2c9473f3c3 100644 --- a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/JdbcParameterBinderTest.java @@ -107,8 +107,8 @@ void bindOrder() throws SQLException { final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root) - .bind(/*parameterIndex=*/ 1, /*columnIndex=*/ 2) - .bind(/*parameterIndex=*/ 2, /*columnIndex=*/ 0) + .bind(/* parameterIndex= */ 1, /* columnIndex= */ 2) + .bind(/* parameterIndex= */ 2, /* columnIndex= */ 0) .build(); assertThat(binder.next()).isFalse(); @@ -166,7 +166,7 @@ void customBinder() throws SQLException { final JdbcParameterBinder binder = JdbcParameterBinder.builder(statement, root) .bind( - /*parameterIndex=*/ 1, + /* parameterIndex= */ 1, new ColumnBinder() { private final IntVector vector = (IntVector) root.getVector(0); 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/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java index 4993a8b1ae..1708c8c12c 100644 --- a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/UnreliableMetaDataTest.java @@ -187,10 +187,11 @@ public void testIncorrectNullability(boolean reuseVectorSchemaRoot) throws Excep final Schema notNullSchema = new Schema( Collections.singletonList( - Field.notNullable(/*name=*/ null, new ArrowType.Int(32, true)))); + Field.notNullable(/* name= */ null, new ArrowType.Int(32, true)))); final Schema nullSchema = new Schema( - Collections.singletonList(Field.nullable(/*name=*/ null, new ArrowType.Int(32, true)))); + Collections.singletonList( + Field.nullable(/* name= */ null, new ArrowType.Int(32, true)))); try (final ResultSet rs = resultSetBuilder.build()) { JdbcToArrowConfig config = diff --git a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumerTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumerTest.java index b1e253794d..bb836578e2 100644 --- a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumerTest.java +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumerTest.java @@ -22,6 +22,7 @@ import java.io.ByteArrayInputStream; import java.io.IOException; +import java.io.InputStream; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.VarBinaryVector; import org.junit.jupiter.api.Test; @@ -65,7 +66,11 @@ public void testConsumeInputStream(byte[][] values, boolean nullable) throws IOE nullable, binaryConsumer -> { for (byte[] value : values) { - binaryConsumer.consume(new ByteArrayInputStream(value)); + if (value != null) { + binaryConsumer.consume(new ByteArrayInputStream(value)); + } else { + binaryConsumer.consume((InputStream) null); + } binaryConsumer.moveWriterPosition(); } }, @@ -119,5 +124,9 @@ public void testConsumeInputStream() throws IOException { testRecords[i] = createBytes(DEFAULT_RECORD_BYTE_COUNT); } testConsumeInputStream(testRecords, false); + + byte[] bytes1 = new byte[] {1, 2, 3}; + byte[] bytes2 = new byte[] {4, 5, 6}; + testConsumeInputStream(new byte[][] {bytes1, null, bytes2}, true); } } diff --git a/adapter/orc/pom.xml b/adapter/orc/pom.xml index 6061feb4ad..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 - 1.9.5 + 2.3.0 test diff --git a/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/OrcReaderTest.java b/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/OrcReaderTest.java index f8eb91a1cc..f48e6bb95e 100644 --- a/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/OrcReaderTest.java +++ b/adapter/orc/src/test/java/org/apache/arrow/adapter/orc/OrcReaderTest.java @@ -38,6 +38,7 @@ import org.apache.orc.TypeDescription; import org.apache.orc.Writer; import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -53,6 +54,7 @@ public static void beforeClass() { allocator = new RootAllocator(MAX_ALLOCATION); } + @Disabled("ORC is flaky: https://github.com/apache/arrow-java/pull/449") @Test public void testOrcJniReader() throws Exception { TypeDescription schema = TypeDescription.fromString("struct"); 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/arrow-variant/src/main/java/module-info.java b/arrow-variant/src/main/java/module-info.java new file mode 100644 index 0000000000..da94173969 --- /dev/null +++ b/arrow-variant/src/main/java/module-info.java @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +@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; + + 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/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.java new file mode 100644 index 0000000000..e3947ac439 --- /dev/null +++ b/arrow-variant/src/main/java/org/apache/arrow/variant/holders/VariantHolder.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 VariantHolder extends ExtensionHolder { + + 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 String toString() { + throw new UnsupportedOperationException(); + } + + @Override + 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 ccb70d5fb3..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 @@ -66,24 +66,24 @@ under the License. - scm:git:https://github.com/apache/arrow.git - scm:git:https://github.com/apache/arrow.git + scm:git:https://github.com/apache/arrow-java.git + scm:git:https://github.com/apache/arrow-java.git main - https://github.com/apache/arrow/tree/${project.scm.tag} + https://github.com/apache/arrow-java/tree/${project.scm.tag} GitHub - https://github.com/apache/arrow/issues + https://github.com/apache/arrow-java/issues - + - 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.30.0 + 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/BufferImportTypeVisitor.java b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java index 2661c12cda..10f690fc87 100644 --- a/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java +++ b/c/src/main/java/org/apache/arrow/c/BufferImportTypeVisitor.java @@ -136,7 +136,7 @@ private ArrowBuf maybeImportBitmap(ArrowType type) { if (buffers[0] == NULL) { return null; } - return importFixedBits(type, 0, /*bitsPerSlot=*/ 1); + return importFixedBits(type, 0, /* bitsPerSlot= */ 1); } @Override @@ -205,7 +205,7 @@ public List visit(ArrowType.FloatingPoint type) { switch (type.getPrecision()) { case HALF: return Arrays.asList( - maybeImportBitmap(type), importFixedBytes(type, 1, /*bytesPerSlot=*/ 2)); + maybeImportBitmap(type), importFixedBytes(type, 1, /* bytesPerSlot= */ 2)); case SINGLE: return Arrays.asList( maybeImportBitmap(type), importFixedBytes(type, 1, Float4Vector.TYPE_WIDTH)); @@ -228,9 +228,8 @@ public List visit(ArrowType.Utf8 type) { type, start, end); - final int len = end - start; offsets.getReferenceManager().retain(); - return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, len)); + return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @@ -279,9 +278,8 @@ public List visit(ArrowType.LargeUtf8 type) { type, start, end); - final long len = end - start; offsets.getReferenceManager().retain(); - return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, len)); + return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @@ -296,9 +294,8 @@ public List visit(ArrowType.Binary type) { type, start, end); - final int len = end - start; offsets.getReferenceManager().retain(); - return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, len)); + return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @@ -320,9 +317,8 @@ public List visit(ArrowType.LargeBinary type) { type, start, end); - final long len = end - start; offsets.getReferenceManager().retain(); - return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, len)); + return Arrays.asList(maybeImportBitmap(type), offsets, importData(type, end)); } } @@ -333,7 +329,7 @@ public List visit(ArrowType.FixedSizeBinary type) { @Override public List visit(ArrowType.Bool type) { - return Arrays.asList(maybeImportBitmap(type), importFixedBits(type, 1, /*bitsPerSlot=*/ 1)); + return Arrays.asList(maybeImportBitmap(type), importFixedBits(type, 1, /* bitsPerSlot= */ 1)); } @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 67ab282de5..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,14 +35,14 @@ 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; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; 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; @@ -74,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; @@ -92,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; @@ -100,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; @@ -181,6 +180,13 @@ boolean roundtrip(FieldVector vector, Class clazz) { clazz.isInstance(imported), String.format("expected %s but was %s", clazz, imported.getClass())); result = VectorEqualsVisitor.vectorEquals(vector, imported); + + if (imported instanceof BaseVariableWidthVector + || imported instanceof BaseLargeVariableWidthVector) { + ArrowBuf offsetBuffer = imported.getOffsetBuffer(); + assertTrue(offsetBuffer.capacity() > 0); + assertEquals(0, offsetBuffer.getInt(0)); + } } // Check that the ref counts of the buffers are the same after the roundtrip @@ -602,6 +608,13 @@ public void testVarCharVector() { } } + @Test + public void testEmptyVarCharVector() { + try (final VarCharVector vector = new VarCharVector("v", allocator)) { + assertTrue(roundtrip(vector, VarCharVector.class)); + } + } + @Test public void testLargeVarBinaryVector() { try (final LargeVarBinaryVector vector = new LargeVarBinaryVector("", allocator)) { @@ -635,6 +648,13 @@ public void testLargeVarCharVector() { } } + @Test + public void testEmptyLargeVarCharVector() { + try (final LargeVarCharVector vector = new LargeVarCharVector("v", allocator)) { + assertTrue(roundtrip(vector, LargeVarCharVector.class)); + } + } + @Test public void testListVector() { try (final ListVector vector = ListVector.empty("v", allocator)) { @@ -789,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(); @@ -809,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()); @@ -935,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); @@ -945,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 @@ -979,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 @@ -1002,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); @@ -1024,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/c/src/test/java/org/apache/arrow/c/StreamTest.java b/c/src/test/java/org/apache/arrow/c/StreamTest.java index 95363fcc32..3dc370c424 100644 --- a/c/src/test/java/org/apache/arrow/c/StreamTest.java +++ b/c/src/test/java/org/apache/arrow/c/StreamTest.java @@ -229,7 +229,7 @@ public void roundtripDictionary() throws Exception { Collections.singletonList( new Field( "dict", - new FieldType(/*nullable=*/ true, indexType, encoding), + new FieldType(/* nullable= */ true, indexType, encoding), Collections.emptyList()))); final List batches = new ArrayList<>(); try (final CDataDictionaryProvider provider = new CDataDictionaryProvider(); @@ -362,7 +362,8 @@ void roundtrip(Schema schema, List batches) throws Exception { private static void assertVectorsEqual(FieldVector expected, FieldVector actual) { assertThat(actual.getField().getType()).isEqualTo(expected.getField().getType()); assertThat(actual.getValueCount()).isEqualTo(expected.getValueCount()); - final Range range = new Range(/*leftStart=*/ 0, /*rightStart=*/ 0, expected.getValueCount()); + final Range range = + new Range(/* leftStart= */ 0, /* rightStart= */ 0, expected.getValueCount()); assertThat(new RangeEqualsVisitor(expected, actual).rangeEquals(range)) .as("Vectors were not equal.\nExpected: %s\nGot: %s", expected, actual) .isTrue(); diff --git a/ci/docker/conda-jni.dockerfile b/ci/docker/conda-jni.dockerfile new file mode 100644 index 0000000000..3f31b74052 --- /dev/null +++ b/ci/docker/conda-jni.dockerfile @@ -0,0 +1,30 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +FROM ghcr.io/mamba-org/micromamba:ubuntu24.04 + +ARG jdk=17 +ARG maven=3.9.9 + +RUN micromamba install -y \ + -c conda-forge \ + cmake \ + compilers \ + maven=${maven} \ + ninja \ + openjdk=${jdk} && \ + micromamba clean --all diff --git a/ci/docker/vcpkg-jni.dockerfile b/ci/docker/vcpkg-jni.dockerfile new file mode 100644 index 0000000000..f2f5d0d45a --- /dev/null +++ b/ci/docker/vcpkg-jni.dockerfile @@ -0,0 +1,32 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +ARG base +FROM ${base} + +# Install Java +# We need Java for JNI headers, but we don't invoke Maven in this build. +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 \ + ARROW_JAVA_CDATA=ON \ + ARROW_JAVA_JNI=ON \ + ARROW_USE_CCACHE=ON + +LABEL org.opencontainers.image.source https://github.com/apache/arrow-java diff --git a/ci/scripts/java_build.sh b/ci/scripts/build.sh similarity index 65% rename from ci/scripts/java_build.sh rename to ci/scripts/build.sh index b5a12d9171..146a40cf7a 100755 --- a/ci/scripts/java_build.sh +++ b/ci/scripts/build.sh @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information @@ -16,22 +17,19 @@ # specific language governing permissions and limitations # under the License. -set -eo pipefail +set -euo pipefail if [[ "${ARROW_JAVA_BUILD:-ON}" != "ON" ]]; then exit fi -arrow_dir=${1} source_dir=${1} build_dir=${2} java_jni_dist_dir=${3} -: ${BUILD_DOCS_JAVA:=OFF} - mvn="mvn -B -DskipTests -Drat.skip=true -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn" -if [ $ARROW_JAVA_SKIP_GIT_PLUGIN ]; then +if [ "${ARROW_JAVA_SKIP_GIT_PLUGIN:-OFF}" = "ON" ]; then mvn="${mvn} -Dmaven.gitcommitid.skip=true" fi @@ -50,38 +48,51 @@ cp -r "${source_dir}/dev" "${build_dir}" poms=$(find "${source_dir}" -not \( -path "${source_dir}"/build -prune \) -type f -name pom.xml) if [[ "$OSTYPE" == "darwin"* ]]; then - poms=$(echo "$poms" | xargs -n1 python -c "import sys; import os.path; print(os.path.relpath(sys.argv[1], '${source_dir}'))") + poms=$(echo "$poms" | xargs -n1 python -c "import sys; import os.path; print(os.path.relpath(sys.argv[1], '${source_dir}'))") else - poms=$(echo "$poms" | xargs -n1 realpath -s --relative-to="${source_dir}") + poms=$(echo "$poms" | xargs -n1 realpath -s --relative-to="${source_dir}") fi for source_root in $(echo "${poms}" | awk -F/ '{print $1}' | sort -u); do - cp -r "${source_dir}/${source_root}" "${build_dir}" + cp -r "${source_dir}/${source_root}" "${build_dir}" done pushd "${build_dir}" -if [ "${ARROW_JAVA_SHADE_FLATBUFFERS}" == "ON" ]; then +# TODO: ARROW_JAVA_SHADE_FLATBUFFERS isn't used for our artifacts. Do +# we need this? +# See also: +# * https://github.com/apache/arrow/issues/22021 +# * https://github.com/apache/arrow-java/issues/67 +if [ "${ARROW_JAVA_SHADE_FLATBUFFERS:-OFF}" == "ON" ]; then mvn="${mvn} -Pshade-flatbuffers" fi -if [ "${ARROW_JAVA_CDATA}" = "ON" ]; then +if [ "${ARROW_JAVA_CDATA:-OFF}" = "ON" ]; then mvn="${mvn} -Darrow.c.jni.dist.dir=${java_jni_dist_dir} -Parrow-c-data" fi -if [ "${ARROW_JAVA_JNI}" = "ON" ]; then +if [ "${ARROW_JAVA_JNI:-OFF}" = "ON" ]; then mvn="${mvn} -Darrow.cpp.build.dir=${java_jni_dist_dir} -Parrow-jni" fi # Use `2 * ncores` threads ${mvn} -T 2C clean install -if [ "${BUILD_DOCS_JAVA}" == "ON" ]; then - # HTTP pooling is turned of to avoid download issues https://issues.apache.org/jira/browse/ARROW-11633 - # GH-43378: Maven site plugins not compatible with multithreading - mkdir -p ${build_dir}/docs/java/reference - ${mvn} -Dcheckstyle.skip=true -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false clean install site - rsync -a target/site/apidocs/ ${build_dir}/docs/java/reference +if [ "${ARROW_JAVA_BUILD_DOCS:-OFF}" == "ON" ]; then + # HTTP pooling is turned off to avoid download issues: + # https://github.com/apache/arrow/issues/27496 + # + # Maven site plugins not compatible with multithreading: + # https://github.com/apache/arrow/issues/43378 + ${mvn} \ + -Dcheckstyle.skip=true \ + -Dhttp.keepAlive=false \ + -Dmaven.wagon.http.pool=false \ + site + rm -rf docs/reference + mkdir -p docs + cp -a target/site/apidocs/ docs/reference fi popd diff --git a/ci/scripts/download_cpp.sh b/ci/scripts/download_cpp.sh new file mode 100755 index 0000000000..3721f62ad5 --- /dev/null +++ b/ci/scripts/download_cpp.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -euxo pipefail + +if [ $# -eq 1 ]; then + version="${1}" +else + version="latest-release" +fi + +url="" + +if [ "${version}" = "latest-release" ]; then + version=$(curl \ + https://raw.githubusercontent.com/apache/arrow-site/refs/heads/main/_data/versions.yml | + grep '^ number:' | + sed -E -e "s/^ number: '|'$//g") +elif [ "${version}" = "latest-rc" ]; then + rc_archive_name=$(curl \ + https://dist.apache.org/repos/dist/dev/arrow/ | + grep -E -o 'apache-arrow-[0-9]+\.[0-9]+\.[0-9]+\-rc[0-9]' | + sort | + uniq | + tail -n1) + rc_version="${rc_archive_name#apache-arrow-}" + version="${rc_version%-rc*}" + url="https://dist.apache.org/repos/dist/dev/arrow/apache-arrow-${rc_version}/apache-arrow-${version}.tar.gz" +fi + +if [ -z "${url}" ]; then + url="https://www.apache.org/dyn/closer.lua?action=download&filename=arrow/arrow-${version}/apache-arrow-${version}.tar.gz" +fi +curl --location --output "apache-arrow-${version}.tar.gz" "${url}" +tar xf "apache-arrow-${version}.tar.gz" +mv "apache-arrow-${version}" arrow diff --git a/ci/scripts/jni_build.sh b/ci/scripts/jni_build.sh new file mode 100755 index 0000000000..c000837987 --- /dev/null +++ b/ci/scripts/jni_build.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -euo pipefail + +# shellcheck source=ci/scripts/util_log.sh +. "$(dirname "${0}")/util_log.sh" + +github_actions_group_begin "Prepare arguments" +source_dir=${1} +arrow_install_dir=${2} +build_dir=${3}/java_jni +# The directory where the final binaries will be stored when scripts finish +dist_dir=${4} +prefix_dir="${build_dir}/java-jni" +github_actions_group_end + +github_actions_group_begin "Clear output directories and leftovers" +rm -rf "${build_dir}" +github_actions_group_end + +github_actions_group_begin "Building Arrow Java C Data Interface native library" + +case "$(uname)" in +Linux) + n_jobs=$(nproc) + ;; +Darwin) + n_jobs=$(sysctl -n hw.logicalcpu) + ;; +*) + n_jobs=${NPROC:-1} + ;; +esac + +: "${ARROW_JAVA_BUILD_TESTS:=${ARROW_BUILD_TESTS:-ON}}" +: "${CMAKE_BUILD_TYPE:=release}" +read -ra EXTRA_CMAKE_OPTIONS <<<"${JAVA_JNI_CMAKE_ARGS:-}" +cmake \ + -S "${source_dir}" \ + -B "${build_dir}" \ + -DARROW_JAVA_JNI_ENABLE_DATASET="${ARROW_DATASET:-OFF}" \ + -DARROW_JAVA_JNI_ENABLE_GANDIVA="${ARROW_GANDIVA:-OFF}" \ + -DARROW_JAVA_JNI_ENABLE_ORC="${ARROW_ORC:-OFF}" \ + -DBUILD_TESTING="${ARROW_JAVA_BUILD_TESTS}" \ + -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \ + -DCMAKE_PREFIX_PATH="${arrow_install_dir}" \ + -DCMAKE_INSTALL_PREFIX="${prefix_dir}" \ + -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD:-OFF}" \ + -DProtobuf_USE_STATIC_LIBS=ON \ + -GNinja \ + "${EXTRA_CMAKE_OPTIONS[@]}" +cmake --build "${build_dir}" --verbose +if [ "${ARROW_JAVA_BUILD_TESTS}" = "ON" ]; then + ctest \ + --output-on-failure \ + --parallel "${n_jobs}" \ + --test-dir "${build_dir}" \ + --timeout 300 +fi +cmake --build "${build_dir}" --target install + +github_actions_group_end + +github_actions_group_begin "Copying artifacts" +mkdir -p "${dist_dir}" +# For Windows. *.dll are installed into bin/ on Windows. +if [ -d "${prefix_dir}/bin" ]; then + mv "${prefix_dir}"/bin/* "${dist_dir}"/ +else + mv "${prefix_dir}"/lib/* "${dist_dir}"/ +fi +github_actions_group_end diff --git a/ci/scripts/jni_full_build.sh b/ci/scripts/jni_full_build.sh new file mode 100755 index 0000000000..5d0aee0555 --- /dev/null +++ b/ci/scripts/jni_full_build.sh @@ -0,0 +1,106 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -euo pipefail + +# shellcheck source=ci/scripts/util_log.sh +. "$(dirname "${0}")/util_log.sh" + +github_actions_group_begin "Prepare arguments" +source_dir="$(cd "${1}" && pwd)" +jni_build_dir="$(cd "${2}" && pwd)" +dist_dir="${3}" +rm -rf "${dist_dir}" +mkdir -p "${dist_dir}" +dist_dir="$(cd "${dist_dir}" && pwd)" +github_actions_group_end + +github_actions_group_begin "Clear old artifacts" +# Ensure that there is no old artifacts inside the maven repository +maven_repo=~/.m2/repository/org/apache/arrow +if [ -d "$maven_repo" ]; then + find "$maven_repo" \ + "(" -name "*.jar" -o -name "*.zip" -o -name "*.pom" ")" \ + -exec echo {} ";" \ + -exec rm -rf {} ";" +fi +github_actions_group_end + +github_actions_group_begin "Generate dummy GPG key" +# Generate dummy GPG key for -Papache-release. +# -Papache-release generates signs (*.asc) of artifacts. +# We don't use these signs in our release process. +( + echo "Key-Type: RSA" + echo "Key-Length: 4096" + echo "Name-Real: Build" + echo "Name-Email: build@example.com" + echo "%no-protection" +) | + gpg --full-generate-key --batch +github_actions_group_end + +pushd "${source_dir}" +github_actions_group_begin "Build .jar" +# build the entire project +mvn \ + --no-transfer-progress \ + -Darrow.c.jni.dist.dir="${jni_build_dir}" \ + -Darrow.cpp.build.dir="${jni_build_dir}" \ + -Papache-release \ + -Parrow-c-data \ + -Parrow-jni \ + clean \ + install +github_actions_group_end +github_actions_group_begin "Build docs" +# build docs +mvn \ + --no-transfer-progress \ + -Darrow.c.jni.dist.dir="${jni_build_dir}" \ + -Darrow.cpp.build.dir="${jni_build_dir}" \ + -Dcheckstyle.skip=true \ + -Dhttp.keepAlive=false \ + -Dmaven.wagon.http.pool=false \ + -Parrow-c-data \ + -Parrow-jni \ + site +github_actions_group_end +popd + +github_actions_group_begin "Prepare artifacts" +# copy all jar, zip and pom files to the distribution folder +find ~/.m2/repository/org/apache/arrow \ + "(" \ + -name "*.jar" -o \ + -name "*.json" -o \ + -name "*.pom" -o \ + -name "*.xml" -o \ + -name "*.zip" \ + ")" \ + -exec echo "{}" ";" \ + -exec cp "{}" "${dist_dir}" ";" + +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 new file mode 100755 index 0000000000..65ab450666 --- /dev/null +++ b/ci/scripts/jni_macos_build.sh @@ -0,0 +1,117 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# This script is like java_jni_build.sh, but is meant for release artifacts +# and hardcodes assumptions about the environment it is being run in. + +set -euo pipefail + +# shellcheck source=ci/scripts/util_log.sh +. "$(dirname "${0}")/util_log.sh" + +github_actions_group_begin "Prepare arguments" +source_dir="$(cd "${1}" && pwd)" +arrow_dir="$(cd "${2}" && pwd)" +build_dir="${3}" +normalized_arch="$(arch)" +case "${normalized_arch}" in +arm64) + normalized_arch=aarch_64 + ;; +i386) + normalized_arch=x86_64 + ;; +esac +# The directory where the final binaries will be stored when scripts finish +dist_dir="${4}" +github_actions_group_end + +github_actions_group_begin "Clear output directories and leftovers" +rm -rf "${build_dir}" +rm -rf "${dist_dir}" + +mkdir -p "${build_dir}" +build_dir="$(cd "${build_dir}" && pwd)" +github_actions_group_end + +: "${ARROW_USE_CCACHE:=ON}" +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics before build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +fi + +github_actions_group_begin "Building Arrow C++ libraries" +install_dir="${build_dir}/cpp-install" + +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" \ + --preset=ninja-release-jni-macos \ + -DCMAKE_INSTALL_PREFIX="${install_dir}" +cmake --build "${build_dir}/cpp" --target install +github_actions_group_end + +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}" \ + "${build_dir}" \ + "${dist_dir}" + +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics after build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +fi + +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 \ + --allow libarrow_dataset_jni \ + --allow libarrow_orc_jni \ + --allow libc++ \ + --allow libcurl \ + --allow libgandiva_jni \ + --allow libncurses \ + --allow libobjc \ + --allow libz \ + --allow libz3 \ + "arrow_cdata_jni/${normalized_arch}/libarrow_cdata_jni.dylib" \ + "arrow_dataset_jni/${normalized_arch}/libarrow_dataset_jni.dylib" \ + "arrow_orc_jni/${normalized_arch}/libarrow_orc_jni.dylib" \ + "gandiva_jni/${normalized_arch}/libgandiva_jni.dylib" +popd +github_actions_group_end diff --git a/ci/scripts/jni_manylinux_build.sh b/ci/scripts/jni_manylinux_build.sh new file mode 100755 index 0000000000..3577c37ab3 --- /dev/null +++ b/ci/scripts/jni_manylinux_build.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# This script is like java_jni_build.sh, but is meant for release artifacts +# and hardcodes assumptions about the environment it is being run in. + +set -euo pipefail + +# shellcheck source=ci/scripts/util_log.sh +. "$(dirname "${0}")/util_log.sh" + +github_actions_group_begin "Prepare arguments" +source_dir="$(cd "${1}" && pwd)" +arrow_dir="$(cd "${2}" && pwd)" +build_dir="${3}" +# The directory where the final binaries will be stored when scripts finish +dist_dir="${4}" +github_actions_group_end + +github_actions_group_begin "Install Archery" +pip install -e "${arrow_dir}/dev/archery[all]" +github_actions_group_end + +github_actions_group_begin "Clear output directories and leftovers" +rm -rf "${build_dir}" +rm -rf "${dist_dir}" + +mkdir -p "${build_dir}" +build_dir="$(cd "${build_dir}" && pwd)" +github_actions_group_end + +: "${ARROW_USE_CCACHE:=ON}" +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics before build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +fi + +github_actions_group_begin "Building Arrow C++ libraries" + +: "${VCPKG_ROOT:=/opt/vcpkg}" +: "${VCPKG_FEATURE_FLAGS:=-manifests}" +: "${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 AWS_EC2_METADATA_DISABLED=TRUE + +install_dir="${build_dir}/cpp-install" + +cmake \ + -S "${arrow_dir}/cpp" \ + -B "${build_dir}/cpp" \ + --preset=ninja-release-jni-linux \ + -DCMAKE_INSTALL_PREFIX="${install_dir}" +cmake --build "${build_dir}/cpp" +cmake --install "${build_dir}/cpp" +github_actions_group_end + +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 +"${source_dir}/ci/scripts/jni_build.sh" \ + "${source_dir}" \ + "${install_dir}" \ + "${build_dir}" \ + "${dist_dir}" + +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics after build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +fi + +github_actions_group_begin "Checking shared dependencies for libraries" +normalized_arch="$(arch)" +case "${normalized_arch}" in +aarch64) + normalized_arch=aarch_64 + ;; +esac +pushd "${dist_dir}" +archery linking check-dependencies \ + --allow ld-linux-aarch64 \ + --allow ld-linux-x86-64 \ + --allow libc \ + --allow libdl \ + --allow libgcc_s \ + --allow libm \ + --allow libpthread \ + --allow librt \ + --allow libstdc++ \ + --allow libz \ + --allow linux-vdso \ + arrow_cdata_jni/"${normalized_arch}"/libarrow_cdata_jni.so \ + arrow_dataset_jni/"${normalized_arch}"/libarrow_dataset_jni.so \ + arrow_orc_jni/"${normalized_arch}"/libarrow_orc_jni.so \ + gandiva_jni/"${normalized_arch}"/libgandiva_jni.so +popd +github_actions_group_end diff --git a/ci/scripts/jni_windows_build.sh b/ci/scripts/jni_windows_build.sh new file mode 100755 index 0000000000..6503ac63e5 --- /dev/null +++ b/ci/scripts/jni_windows_build.sh @@ -0,0 +1,139 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -euo pipefail + +# shellcheck source=ci/scripts/util_log.sh +. "$(dirname "${0}")/util_log.sh" + +github_actions_group_begin "Prepare arguments" +source_dir="$(cd "${1}" && pwd)" +arrow_dir="$(cd "${2}" && pwd)" +build_dir="${3}" +# The directory where the final binaries will be stored when scripts finish +dist_dir="${4}" +github_actions_group_end + +github_actions_group_begin "Clear output directories and leftovers" +rm -rf "${build_dir}" + +mkdir -p "${build_dir}" +build_dir="$(cd "${build_dir}" && pwd)" +github_actions_group_end + +: "${ARROW_USE_CCACHE:=ON}" +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics before build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +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}" # TODO: ON +export ARROW_BUILD_TESTS +: "${ARROW_DATASET:=ON}" +export ARROW_DATASET +: "${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 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_ORC="${ARROW_ORC}" \ + -DARROW_PARQUET="${ARROW_PARQUET}" \ + -DARROW_S3="${ARROW_S3}" \ + -DARROW_USE_CCACHE="${ARROW_USE_CCACHE}" \ + -DARROW_WITH_BROTLI=ON \ + -DARROW_WITH_LZ4=ON \ + -DARROW_WITH_SNAPPY=ON \ + -DARROW_WITH_ZSTD=ON \ + -DCMAKE_BUILD_TYPE="${CMAKE_BUILD_TYPE}" \ + -DCMAKE_INSTALL_PREFIX="${install_dir}" \ + -DCMAKE_UNITY_BUILD="${CMAKE_UNITY_BUILD}" \ + -GNinja +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" + # unstable + exclude_tests="${exclude_tests}|arrow-compute-hash-join-node-test" + exclude_tests="${exclude_tests}|arrow-dataset-scanner-test" + # strptime + exclude_tests="${exclude_tests}|arrow-utility-test" + ctest \ + --exclude-regex "${exclude_tests}" \ + --label-regex unittest \ + --output-on-failure \ + --parallel "$(nproc)" \ + --timeout 300 + github_actions_group_end +fi + +"${source_dir}/ci/scripts/jni_build.sh" \ + "${source_dir}" \ + "${install_dir}" \ + "${build_dir}" \ + "${dist_dir}" + +if [ "${ARROW_USE_CCACHE}" == "ON" ]; then + github_actions_group_begin "ccache statistics after build" + ccache -sv 2>/dev/null || ccache -s + github_actions_group_end +fi + +github_actions_group_begin "Checking shared dependencies for libraries" +normalized_arch="$(arch)" +case "${normalized_arch}" in +aarch64) + normalized_arch=aarch_64 + ;; +esac +pushd "${dist_dir}" +# TODO +# archery linking check-dependencies \ +# --allow libm \ +# --allow librt \ +# --allow libz \ +# arrow_cdata_jni/"${normalized_arch}"/libarrow_cdata_jni.dll \ +# arrow_dataset_jni/"${normalized_arch}"/libarrow_dataset_jni.dll \ +# arrow_orc_jni/"${normalized_arch}"/libarrow_orc_jni.dll +popd +github_actions_group_end diff --git a/ci/scripts/java_test.sh b/ci/scripts/test.sh similarity index 72% rename from ci/scripts/java_test.sh rename to ci/scripts/test.sh index 9d4bc018b3..8061ee455d 100755 --- a/ci/scripts/java_test.sh +++ b/ci/scripts/test.sh @@ -1,4 +1,5 @@ #!/usr/bin/env bash +# # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information @@ -22,18 +23,22 @@ if [[ "${ARROW_JAVA_TEST:-ON}" != "ON" ]]; then exit fi -arrow_dir=${1} -source_dir=${1} +source_dir="$(cd "${1}" && pwd)" build_dir=${2} java_jni_dist_dir=${3} +if [ -d "${java_jni_dist_dir}" ]; then + java_jni_dist_dir="$(cd "${java_jni_dist_dir}" && pwd)" +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} +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 @@ -42,14 +47,17 @@ if [ "${ARROW_JAVA_JNI}" = "ON" ]; then projects+=(gandiva) fi if [ "${#projects[@]}" -gt 0 ]; then - ${mvn} clean test \ - -Parrow-jni \ - -pl $(IFS=,; echo "${projects[*]}") \ - -Darrow.cpp.build.dir=${java_jni_dist_dir} + ${mvn} test \ + -Parrow-jni \ + -pl "$( + IFS=, + echo \""${projects[*]}"\" + )" \ + -Darrow.cpp.build.dir="${java_jni_dist_dir}" 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/ci/scripts/util_log.sh b/ci/scripts/util_log.sh new file mode 100644 index 0000000000..c8ee48bbb2 --- /dev/null +++ b/ci/scripts/util_log.sh @@ -0,0 +1,28 @@ +# shellcheck shell=bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +github_actions_group_begin() { + echo "::group::$1" + set -x +} + +github_actions_group_end() { + set +x + echo "::endgroup::" +} diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 0000000000..4fd825e5a5 --- /dev/null +++ b/compose.yaml @@ -0,0 +1,114 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Usage +# ----- +# +# The docker compose file is parametrized using environment variables, the +# defaults are set in .env file. +# +# Example: +# $ ARCH=arm64v8 docker compose build java +# $ ARCH=arm64v8 docker compose run java + +volumes: + ccache-cache: + name: ccache-cache + maven-cache: + name: maven-cache + +services: + ubuntu: + # Build and test arrow-java on Ubuntu. + # + # Usage: + # docker compose build ubuntu + # docker compose run ubuntu + # Parameters: + # MAVEN: 3.9.9 + # JDK: 17, 21 + image: ${ARCH}/maven:${MAVEN}-eclipse-temurin-${JDK} + volumes: + - .:/arrow-java:delegated + - ${DOCKER_VOLUME_PREFIX}maven-cache:/root/.m2:delegated + command: + /bin/bash -c " + /arrow-java/ci/scripts/build.sh /arrow-java /build /jni && + /arrow-java/ci/scripts/test.sh /arrow-java /build /jni" + + conda-jni-cdata: + # Builds and tests just the C Data Interface JNI library and JARs. + # (No dependencies on arrow-cpp.) + # This build isn't meant for distribution. It's for testing only. + # + # Usage: + # docker compose build conda-jni-cdata + # docker compose run conda-jni-cdata + # Parameters: + # MAVEN: 3.9.9 + # JDK: 17, 21 + image: ${REPO}:${ARCH}-conda-java-${JDK}-maven-${MAVEN}-jni-integration + build: + context: . + dockerfile: ci/docker/conda-jni.dockerfile + cache_from: + - ${REPO}:${ARCH}-conda-java-${JDK}-maven-${MAVEN}-jni-integration + args: + jdk: ${JDK} + maven: ${MAVEN} + # required to use micromamba with rootless docker + # https://github.com/mamba-org/micromamba-docker/issues/407#issuecomment-2088523507 + user: root + volumes: + - .:/arrow-java:delegated + - ${DOCKER_VOLUME_PREFIX}maven-cache:/root/.m2:delegated + environment: + ARROW_JAVA_CDATA: "ON" + command: + /bin/bash -c " + /arrow-java/ci/scripts/jni_build.sh /arrow-java /build/jni /build /jni && + /arrow-java/ci/scripts/build.sh /arrow-java /build /jni && + /arrow-java/ci/scripts/test.sh /arrow-java /build /jni" + + vcpkg-jni: + # Builds all the JNI libraries, but not the JARs. + # (Requires arrow-cpp.) + # The artifacts from this build are meant to be used for packaging. + # + # Usage: + # docker compose build vcpkg-jni + # docker compose run vcpkg-jni + image: ${REPO}:${ARCH}-vcpkg-jni-${VCPKG} + build: + context: . + dockerfile: ci/docker/vcpkg-jni.dockerfile + cache_from: + - ${REPO}:${ARCH}-vcpkg-jni-${VCPKG} + args: + base: ${ARROW_REPO}:${ARCH}-cpp-jni-${VCPKG} + volumes: + - .:/arrow-java:delegated + - ${ARROW_REPO_ROOT}:/arrow:delegated + - ${DOCKER_VOLUME_PREFIX}ccache-cache:/ccache:delegated + - ${DOCKER_VOLUME_PREFIX}maven-cache:/root/.m2:delegated + environment: + ARROW_JAVA_CDATA: "ON" + CCACHE_DIR: "/ccache" + command: + ["/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 8cc4909034..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.6-7 + 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 edd52604bc..f268e815fe 100644 --- a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java +++ b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java @@ -41,7 +41,7 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu Integer.MAX_VALUE); byte[] inBytes = new byte[(int) uncompressedBuffer.writerIndex()]; - uncompressedBuffer.getBytes(/*index=*/ 0, inBytes); + uncompressedBuffer.getBytes(/* index= */ 0, inBytes); ByteArrayOutputStream baos = new ByteArrayOutputStream(); try (InputStream in = new ByteArrayInputStream(inBytes); OutputStream out = new FramedLZ4CompressorOutputStream(baos)) { @@ -80,8 +80,15 @@ 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.setBytes(/* index= */ 0, outBytes); decompressedBuffer.writerIndex(decompressedLength); return decompressedBuffer; } 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 6e48aae71f..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,10 +44,10 @@ 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); + /* srcSize= */ uncompressedBuffer.writerIndex(), + /* level= */ this.compressionLevel); if (Zstd.isError(bytesWritten)) { compressedBuffer.close(); throw new RuntimeException("Error compressing: " + Zstd.getErrorName(bytesWritten)); @@ -64,11 +64,12 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu Zstd.decompressUnsafe( uncompressedBuffer.memoryAddress(), decompressedLength, - /*src=*/ compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, + /* src= */ compressedBuffer.memoryAddress() + + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH); if (Zstd.isError(decompressedSize)) { uncompressedBuffer.close(); - throw new RuntimeException("Error decompressing: " + Zstd.getErrorName(decompressedLength)); + throw new RuntimeException("Error decompressing: " + Zstd.getErrorName(decompressedSize)); } if (decompressedLength != decompressedSize) { uncompressedBuffer.close(); diff --git a/compression/src/test/java/org/apache/arrow/compression/TestArrowReaderWriterWithCompression.java b/compression/src/test/java/org/apache/arrow/compression/TestArrowReaderWriterWithCompression.java index d7318e306c..03446e5316 100644 --- a/compression/src/test/java/org/apache/arrow/compression/TestArrowReaderWriterWithCompression.java +++ b/compression/src/test/java/org/apache/arrow/compression/TestArrowReaderWriterWithCompression.java @@ -141,7 +141,7 @@ private Dictionary createDictionary(VarCharVector dictionaryVector) { return new Dictionary( dictionaryVector, - new DictionaryEncoding(/*id=*/ 1L, /*ordered=*/ false, /*indexType=*/ null)); + new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); } @Test 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 21f67f1a69..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.14.4 - 1.12.0 + 1.17.1 + 1.12.1 @@ -130,7 +130,7 @@ under the License. org.apache.orc orc-core - 1.9.5 + 2.3.0 test @@ -156,7 +156,7 @@ under the License. commons-io commons-io - 2.17.0 + 2.22.0 test diff --git a/dataset/src/main/cpp/jni_util.cc b/dataset/src/main/cpp/jni_util.cc index 1fd15696e6..35bfb328f0 100644 --- a/dataset/src/main/cpp/jni_util.cc +++ b/dataset/src/main/cpp/jni_util.cc @@ -187,7 +187,7 @@ ReservationListenableMemoryPool::~ReservationListenableMemoryPool() {} std::string Describe(JNIEnv* env, jthrowable t) { jclass describer_class = env->FindClass("org/apache/arrow/dataset/jni/JniExceptionDescriber"); - DCHECK_NE(describer_class, nullptr); + ARROW_DCHECK_NE(describer_class, nullptr); jmethodID describe_method = env->GetStaticMethodID( describer_class, "describe", "(Ljava/lang/Throwable;)Ljava/lang/String;"); std::string description = JStringToCString( @@ -197,7 +197,7 @@ std::string Describe(JNIEnv* env, jthrowable t) { bool IsErrorInstanceOf(JNIEnv* env, jthrowable t, std::string class_name) { jclass java_class = env->FindClass(class_name.c_str()); - DCHECK_NE(java_class, nullptr) << "Could not find Java class " << class_name; + ARROW_DCHECK_NE(java_class, nullptr) << "Could not find Java class " << class_name; return env->IsInstanceOf(t, java_class); } 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/dataset/src/test/java/org/apache/arrow/dataset/TestAllTypes.java b/dataset/src/test/java/org/apache/arrow/dataset/TestAllTypes.java index eb73663191..79800f3570 100644 --- a/dataset/src/test/java/org/apache/arrow/dataset/TestAllTypes.java +++ b/dataset/src/test/java/org/apache/arrow/dataset/TestAllTypes.java @@ -95,11 +95,9 @@ private VectorSchemaRoot generateAllTypesVector(BufferAllocator allocator) { // DenseUnion List childFields = new ArrayList<>(); childFields.add( - new Field( - "int-child", new FieldType(false, new ArrowType.Int(32, true), null, null), null)); + new Field("int-child", FieldType.notNullable(new ArrowType.Int(32, true)), null)); Field structField = - new Field( - "struct", new FieldType(true, ArrowType.Struct.INSTANCE, null, null), childFields); + new Field("struct", FieldType.nullable(ArrowType.Struct.INSTANCE), childFields); Field[] fields = new Field[] { Field.nullablePrimitive("null", ArrowType.Null.INSTANCE), @@ -239,7 +237,11 @@ private VectorSchemaRoot generateAllTypesVector(BufferAllocator allocator) { largeListWriter.integer().writeInt(1); largeListWriter.endList(); - ((StructVector) root.getVector("struct")).getChild("int-child", IntVector.class).set(1, 1); + IntVector intChildVector = + ((StructVector) root.getVector("struct")).getChild("int-child", IntVector.class); + // Non-nullable vector, make sure to fill all slots + intChildVector.set(0, 0); + intChildVector.set(1, 1); return root; } diff --git a/dataset/src/test/java/org/apache/arrow/dataset/TestDataset.java b/dataset/src/test/java/org/apache/arrow/dataset/TestDataset.java index f3ca04d77b..4b155137ed 100644 --- a/dataset/src/test/java/org/apache/arrow/dataset/TestDataset.java +++ b/dataset/src/test/java/org/apache/arrow/dataset/TestDataset.java @@ -123,8 +123,6 @@ protected void assertParquetFileEquals(String expectedURI, String actualURI) thr VectorSchemaRoot actualVsr = VectorSchemaRoot.create(actualFactory.inspect(), rootAllocator())) { - // fast-fail by comparing metadata - assertEquals(expectedBatches.toString(), actualBatches.toString()); // compare ArrowRecordBatches assertEquals(expectedBatches.size(), actualBatches.size()); VectorLoader expectLoader = new VectorLoader(expectVsr); diff --git a/dev/checkstyle/suppressions.xml b/dev/checkstyle/suppressions.xml index e8669c54e6..b1841d7fe6 100644 --- a/dev/checkstyle/suppressions.xml +++ b/dev/checkstyle/suppressions.xml @@ -37,7 +37,7 @@ - + diff --git a/docker-compose.yml b/dev/release/.env.example similarity index 50% rename from docker-compose.yml rename to dev/release/.env.example index 103f2f3ad0..4a57e34ed3 100644 --- a/docker-compose.yml +++ b/dev/release/.env.example @@ -15,33 +15,18 @@ # specific language governing permissions and limitations # under the License. -# Usage -# ----- +# The GitHub token to upload artifacts to GitHub Release. # -# The docker compose file is parametrized using environment variables, the -# defaults are set in .env file. -# -# Example: -# $ ARCH=arm64v8 docker compose build java -# $ ARCH=arm64v8 docker compose run java - -volumes: - maven-cache: - name: maven-cache +# You must set this. +#GH_TOKEN=secret -services: - java: - # Usage: - # docker compose build java - # docker compose run java - # Parameters: - # MAVEN: 3.9.6 - # JDK: 11, 17, 21 - image: ${ARCH}/maven:${MAVEN}-eclipse-temurin-${JDK} - volumes: &java-volumes - - .:/arrow-java:delegated - - ${DOCKER_VOLUME_PREFIX}maven-cache:/root/.m2:delegated - command: &java-command > - /bin/bash -c " - /arrow-java/ci/scripts/java_build.sh /arrow-java /build && - /arrow-java/ci/scripts/java_test.sh /arrow-java /build" +# The GPG key ID to sign artifacts. The GPG key ID must be registered +# to both of the followings: +# +# * https://dist.apache.org/repos/dist/dev/arrow/KEYS +# * https://dist.apache.org/repos/dist/release/arrow/KEYS +# +# See these files how to import your GPG key ID to these files. +# +# You must set this. +#GPG_KEY_ID=08D3564B7C6A9CAFBFF6A66791D18FCF079F8007 diff --git a/dev/release/README.md b/dev/release/README.md new file mode 100644 index 0000000000..8aee0fd106 --- /dev/null +++ b/dev/release/README.md @@ -0,0 +1,293 @@ + + +# Release + +## Overview + + 1. Test the revision to be released + 2. Bump version for new release (detailed later) + 3. Prepare RC and vote (detailed later) + 4. Publish (detailed later) + 5. Bump version for new development (detailed later) + +### Prepare release environment + +This step is needed only when you act as a release manager first time. + +We use the following variables in multiple steps: + +* `GH_TOKEN`: GitHub personal access token to automate GitHub related + operations +* `GPG_KEY_ID`: PGP key ID that is used for signing official artifacts + by GnuPG + +We use `dev/release/.env` to share these variables in multiple +steps. You can use `dev/release/.env.example` as a template: + +```console +$ cp dev/release/.env{.example,} +$ chmod go-r dev/release/.env +$ editor dev/release/.env +``` + +See +https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens +how to prepare GitHub personal access token for `GH_TOKEN`. + +Note that you also need to install `gh` command because our scripts +use `gh` command to use GitHub API. See +https://github.com/cli/cli#installation how to install `gh` +command. + +If you don't have a PGP key for `GPG_KEY_ID`, see +https://infra.apache.org/release-signing.html#genegrate how to +generate your PGP key. + +Your PGP key must be registered to the followings: + + * https://dist.apache.org/repos/dist/dev/arrow/KEYS + * https://dist.apache.org/repos/dist/release/arrow/KEYS + +See the header comment of them how to add a PGP key. + +Apache arrow committers can update them by Subversion client with +their ASF account. e.g.: + +```console +$ svn co https://dist.apache.org/repos/dist/dev/arrow +$ cd arrow +$ head KEYS +(This shows how to update KEYS) +$ svn ci KEYS +``` + +### Bump version for new release + +Run `dev/release/bump_version.sh` on a working copy of your fork not +`git@github.com:apache/arrow-java`: + +```console +$ git clone git@github.com:${YOUR_GITHUB_ACCOUNT}/arrow-java.git arrow-java.${YOUR_GITHUB_ACCOUNT} +$ cd arrow-java.${YOUR_GITHUB_ACCOUNT} +$ dev/release/bump_version.sh ${NEW_VERSION} +``` + +Here is an example to bump version to 19.0.0: + +``` +$ dev/release/bump_version.sh 19.0.0 +``` + +It creates a feature branch and adds a commit that bumps version. This +opens a pull request from the feature branch. + +We need to merge the pull request before we cut a RC. If we try +cutting a RC without merging the pull request, the script to cut a RC +is failed. + +### Prepare RC and vote + +You can use `dev/release/release_rc.sh`. + +Requirements to run `release_rc.sh`: + + * You must be an Apache Arrow committer or PMC member + * You must prepare your PGP key for signing + * You must configure Maven + +Configure Maven to publish artifacts to Apache repositories. You will +need to setup a master password at `~/.m2/settings-security.xml` and +`~/.m2/settings.xml` as specified on [the Apache +guide](https://infra.apache.org/publishing-maven-artifacts.html). It +can be tested with the following command: + +```bash +# You might need to export GPG_TTY=$(tty) to properly prompt for a passphrase +mvn clean install -Papache-release +``` + +Run `dev/release/release_rc.sh` on a working copy of +`git@github.com:apache/arrow-java` not your fork: + +```console +$ git clone git@github.com:apache/arrow-java.git +$ cd arrow-java +$ dev/release/release_rc.sh ${RC} +(Send a vote email to dev@arrow.apache.org. + You can use a draft shown by release_rc.sh for the email.) +``` + +Here is an example to release RC1: + +```console +$ dev/release/release_rc.sh 1 +``` + +The argument of `release_rc.sh` is the RC number. If RC1 has a +problem, we'll increment the RC number such as RC2, RC3 and so on. + +### Publish + +We need to do the followings to publish a new release: + + * Publish the source archive to apache.org + * Publish the binary artifacts to repository.apache.org + +Run `dev/release/release.sh` on a working copy of +`git@github.com:apache/arrow-java` not your fork to publish the source +archive to apache.org: + +```console +$ dev/release/release.sh ${VERSION} ${RC} +``` + +Here is an example to release 19.0.0 RC1: + +```console +$ dev/release/release.sh 19.0.0 1 +``` + +Add the release to ASF's report database via [Apache Committee Report +Helper](https://reporter.apache.org/addrelease.html?arrow). + +You need to do the followings to publish the binary artifacts to +repository.apache.org: + +* Logon to the Apache repository: + https://repository.apache.org/#stagingRepositories +* Select the Arrow staging repository you created for RC: + `orgapachearrow-XXXX` +* Click the `release` button + +### Bump version for new development + +We should bump version in the main branch for new development after we +release a new version. + +Run `dev/release/bump_version.sh` on a working copy of your fork not +`git@github.com:apache/arrow-java`: + +```console +$ git clone git@github.com:${YOUR_GITHUB_ACCOUNT}/arrow-java.git arrow-java.${YOUR_GITHUB_ACCOUNT} +$ cd arrow-java.${YOUR_GITHUB_ACCOUNT} +$ dev/release/bump_version.sh ${NEW_VERSION}-SNAPSHOT +``` + +Here is an example to bump version to 19.0.1-SNAPSHOT: + +``` +$ dev/release/bump_version.sh 19.0.0-SNAPSHOT +``` + +It creates a feature branch and adds a commit that bumps version. This +opens a pull request from the feature branch by `gh pr create`. + +### Close the GitHub milestone + +Close the milestone here, then open a new milestone for the next release: + +https://github.com/apache/arrow-java/milestones + +The milestone should be named after the version (e.g. "18.4.0"). + +### Publish the release blog post + +Open a pull request on +[apache/arrow-site](https://github.com/apache/arrow-site) announcing the new +release and summarizing the changes. Do not include the full changelog, just +important entries (breaking changes, new features, major bug fixes) and a link +to the full changelog. See an [example +PR](https://github.com/apache/arrow-site/pull/594). + +### Announce the new release on the mailing list + +Send an email to "announce@apache.org" from your Apache email, CC'ing +dev@arrow.apache.org/user@arrow.apache.org. See an [example +post](https://lists.apache.org/thread/bxpt0r8kw0ltgywnylqdroskkt6966z4). + +``` +To: announce@apache.org +CC: dev@arrow.apache.org, user@arrow.apache.org +Subject: [ANNOUNCE] Apache Arrow Java 18.2.0 released + +The Apache Arrow community is pleased to announce the Arrow Java 18.2.0 release. + +The release is available now from our website: + https://arrow.apache.org/install/ +and + https://www.apache.org/dyn/closer.cgi/arrow/apache-arrow-java-18.2.0/ + +Read about what's new in the release at: + https://arrow.apache.org/blog/2025/02/19/arrow-java-18.2.0/ + +Read the full changelog: + https://github.com/apache/arrow-java/commits/v18.2.0 + +What is Apache Arrow? +------------------------------- + +Apache Arrow is a universal columnar format and multi-language toolbox +for fast data interchange and in-memory analytics. It houses a set of +canonical in-memory representations of flat and hierarchical data +along with multiple language-bindings for structure manipulation. It +also provides low-overhead streaming and batch messaging, zero-copy +interprocess communication (IPC), and vectorized in-memory analytics +libraries. + +Please report any feedback to the mailing lists: + https://lists.apache.org/list.html?dev@arrow.apache.org + +Regards, +The Apache Arrow community. +``` + +### Announce the new release on social media + +Make a post on our [BlueSky](https://bsky.app/profile/arrow.apache.org) and +[LinkedIn](https://www.linkedin.com/company/apache-arrow/) accounts. (Ask +your fellow PMC members for access if need be, or ask a PMC member to make the +post on your behalf.) The post should link to the blog post. See [example +BlueSky post](https://bsky.app/profile/arrow.apache.org/post/3lioi6ov5h22d) +and [example LinkedIn post](https://www.linkedin.com/posts/apache-arrow_apache-arrow-java-1820-release-activity-7298633716522758144-L71x). + +## Verify + +We have a script to verify a RC. + +You must install the following commands to use the script: + + * `curl` + * `gpg` + * `shasum` or `sha256sum`/`sha512sum` + * `tar` + +To verify a RC, run the following command line: + +```console +$ dev/release/verify_rc.sh ${VERSION} ${RC} +``` + +Here is an example to verify the release 19.0.0 RC1: + +```console +$ dev/release/verify_rc.sh 19.0.0 1 +``` + +If the verification is successful, the message `RC looks good!` is shown. diff --git a/dev/release/bump_version.sh b/dev/release/bump_version.sh new file mode 100755 index 0000000000..68cafb99bd --- /dev/null +++ b/dev/release/bump_version.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -eu + +SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SOURCE_TOP_DIR="$(cd "${SOURCE_DIR}/../../" && pwd)" + +if [ "$#" -ne 1 ]; then + echo "Usage: $0 " + echo " e.g.: $0 19.0.1" + echo " e.g.: $0 20.0.0-SNAPSHOT" + exit 1 +fi + +version=$1 + +if [ ! -f "${SOURCE_DIR}/.env" ]; then + echo "You must create ${SOURCE_DIR}/.env" + echo "You can use ${SOURCE_DIR}/.env.example as template" + exit 1 +fi +. "${SOURCE_DIR}/.env" +export GH_TOKEN + +cd "${SOURCE_TOP_DIR}" + +git_origin_url="$(git remote get-url origin)" +case "${git_origin_url}" in +*apache/arrow-java*) + echo "You must use your fork: ${git_origin_url}" + exit 1 + ;; +*) + : # OK + ;; +esac + +branch="bump-version-${version}" +git switch -c "${branch}" main +mvn versions:set "-DnewVersion=${version}" -DprocessAllModules -DgenerateBackupPoms=false +case "${version}" in +*-SNAPSHOT) + tag=main + ;; +*) + tag=v${version} + ;; +esac +mvn versions:set-scm-tag "-DnewTag=${tag}" -DgenerateBackupPoms=false -pl :arrow-java-root +mvn versions:set-scm-tag "-DnewTag=${tag}" -DgenerateBackupPoms=false -pl :arrow-bom +git add "pom.xml" +git add "**/pom.xml" +git commit -m "MINOR: Bump version to ${version}" +git push --set-upstream origin "${branch}" +gh pr create --fill --repo apache/arrow-java diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index 8efd379a73..0999f1a275 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -16,4 +16,6 @@ # under the License. .gitmodules +.github/pull_request_template.md dataset/src/test/resources/data/student.csv +docs/Makefile diff --git a/dev/release/release.sh b/dev/release/release.sh new file mode 100755 index 0000000000..d1db7ad05a --- /dev/null +++ b/dev/release/release.sh @@ -0,0 +1,106 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +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" + exit 1 +fi + +version=$1 +rc=$2 + +if [ ! -f "${SOURCE_DIR}/.env" ]; then + echo "You must create ${SOURCE_DIR}/.env" + echo "You can use ${SOURCE_DIR}/.env.example as template" + exit 1 +fi +. "${SOURCE_DIR}/.env" +export GH_TOKEN + +git_origin_url="$(git remote get-url origin)" +repository="${git_origin_url#*github.com?}" +repository="${repository%.git}" +case "${git_origin_url}" in +git@github.com:apache/arrow-java.git | https://github.com/apache/arrow-java.git) + : # OK + ;; +*) + echo "This script must be ran with working copy of apache/arrow-java." + echo "The origin's URL: ${git_origin_url}" + exit 1 + ;; +esac + +tag="v${version}" +rc_tag="${tag}-rc${rc}" +echo "Tagging for release: ${tag}" +git tag -a -m "${version}" "${tag}" "${rc_tag}^{}" +git push origin "${tag}" + +release_id="apache-arrow-java-${version}" +source_archive="apache-arrow-java-${version}.tar.gz" +dist_url="https://dist.apache.org/repos/dist/release/arrow" +dist_base_dir="dev/release/dist" +dist_dir="${dist_base_dir}/${release_id}" +echo "Checking out ${dist_url}" +rm -rf "${dist_base_dir}" +svn co --depth=empty "${dist_url}" "${dist_base_dir}" +gh release download "${rc_tag}" \ + --dir "${dist_dir}" \ + --pattern "${source_archive}*" \ + --repo "${repository}" \ + --skip-existing + +echo "Uploading to release/" +pushd "${dist_base_dir}" +svn add "${release_id}" +svn ci -m "Apache Arrow Java ${version}" +popd +rm -rf "${dist_base_dir}" + +echo "Keep only the latest versions" +old_releases=$( + svn ls https://dist.apache.org/repos/dist/release/arrow/ | + grep -E '^apache-arrow-java-' | + sort --version-sort --reverse | + tail -n +2 +) +for old_release_version in ${old_releases}; do + echo "Remove old release ${old_release_version}" + svn \ + delete \ + -m "Remove old Apache Arrow Java release: ${old_release_version}" \ + "https://dist.apache.org/repos/dist/release/arrow/${old_release_version}" +done + +echo "Success! The release is available here:" +echo " https://dist.apache.org/repos/dist/release/arrow/${release_id}" +echo +echo "Add this release to ASF's report database:" +echo " https://reporter.apache.org/addrelease.html?arrow" +echo +echo "Release binary artifacts in repository.apache.org:" +echo "1. Open https://repository.apache.org/#stagingRepositories" +echo "2. Select the repository for RC: orgapachearrow-XXXX" +echo "3. Click the Release button" diff --git a/dev/release/release_rc.sh b/dev/release/release_rc.sh new file mode 100755 index 0000000000..0920edbe35 --- /dev/null +++ b/dev/release/release_rc.sh @@ -0,0 +1,264 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -eu + +SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SOURCE_TOP_DIR="$(cd "${SOURCE_DIR}/../../" && pwd)" + +if [ "$#" -ne 1 ]; then + echo "Usage: $0 " + echo " e.g.: $0 1" + exit 1 +fi + +rc=$1 + +: "${RELEASE_DEFAULT:=1}" +: "${RELEASE_PULL:=${RELEASE_DEFAULT}}" +: "${RELEASE_PUSH_TAG:=${RELEASE_DEFAULT}}" +: "${RELEASE_SIGN:=${RELEASE_DEFAULT}}" +: "${RELEASE_UPLOAD:=${RELEASE_DEFAULT}}" + +if [ ! -f "${SOURCE_DIR}/.env" ]; then + echo "You must create ${SOURCE_DIR}/.env" + echo "You can use ${SOURCE_DIR}/.env.example as template" + exit 1 +fi +. "${SOURCE_DIR}/.env" +export GH_TOKEN + +cd "${SOURCE_TOP_DIR}" + +if [ "${RELEASE_PULL}" -gt 0 ] || [ "${RELEASE_PUSH_TAG}" -gt 0 ]; then + git_origin_url="$(git remote get-url origin)" + case "${git_origin_url}" in + git@github.com:apache/arrow-java.git | https://github.com/apache/arrow-java.git) + : # OK + ;; + *) + echo "This script must be ran with working copy of apache/arrow-java." + echo "The origin's URL: ${git_origin_url}" + exit 1 + ;; + esac +fi + +if [ "${RELEASE_PULL}" -gt 0 ]; then + echo "Ensure using the latest commit" + git checkout main + git pull --ff-only +fi + +version=$(grep -o '^ .*' "pom.xml" | + sed \ + -e 's,^ ,,' \ + -e 's,$,,') + +case "${version}" in +*-SNAPSHOT) + echo "Version isn't bumped: ${version}" + echo "Run dev/release/bump_version.sh before you run this script." + exit 1 + ;; +esac + +rc_tag="v${version}-rc${rc}" +if [ "${RELEASE_PUSH_TAG}" -gt 0 ]; then + echo "Tagging for RC: ${rc_tag}" + git tag -a -m "${version} RC${rc}" "${rc_tag}" + git push origin "${rc_tag}" +fi + +rc_hash="$(git rev-list --max-count=1 "${rc_tag}")" + +artifacts_dir="apache-arrow-java-${version}-rc${rc}" +signed_artifacts_dir="${artifacts_dir}-signed" + +if [ "${RELEASE_SIGN}" -gt 0 ]; then + git_origin_url="$(git remote get-url origin)" + repository="${git_origin_url#*github.com?}" + repository="${repository%.git}" + + echo "Looking for GitHub Actions workflow on ${repository}:${rc_tag}" + run_id="" + while true; do + echo "Waiting for run to start..." + run_id=$(gh run list \ + --branch "${rc_tag}" \ + --jq ".[].databaseId" \ + --json 'databaseId' \ + --limit 1 \ + --repo "${repository}" \ + --workflow rc.yml) + if [ -n "${run_id}" ]; then + break + fi + sleep 600 + done + + echo "Found GitHub Actions workflow with ID: ${run_id}" + gh run watch \ + --exit-status "${run_id}" \ + --interval 600 \ + --repo "${repository}" + + echo "Downloading artifacts from GitHub Releases" + gh release download "${rc_tag}" \ + --dir "${artifacts_dir}" \ + --repo "${repository}" \ + --skip-existing + + echo "Signing artifacts" + rm -rf "${signed_artifacts_dir}" + mkdir -p "${signed_artifacts_dir}" + for artifact in "${artifacts_dir}"/*; do + case "${artifact}" in + *.asc | *.sha256 | *.sha512) + continue + ;; + esac + gpg \ + --armor \ + --detach-sig \ + --local-user "${GPG_KEY_ID}" \ + --output "${signed_artifacts_dir}/$(basename "${artifact}").asc" \ + "${artifact}" + done +fi + +# arrow-c-data-18.2.0-sources.jar -> +# jar +extract_type() { + local path="$1" + echo "${path}" | grep -o "[^.]*$" +} + +# arrow-c-data-18.2.0-sources.jar arrow-c-data-18.2.0 -> +# sources +extract_classifier() { + local path="$1" + local base="$2" + basename "${path}" | sed -e "s/^${base}-//g" -e "s/\.[^.]*$//g" +} + +if [ "${RELEASE_UPLOAD}" -gt 0 ]; then + echo "Uploading signature" + gh release upload "${rc_tag}" \ + --clobber \ + --repo "${repository}" \ + "${signed_artifacts_dir}"/*.asc + + echo "Uploading packages" + for pom in "${artifacts_dir}"/*.pom; do + base=$(basename "${pom}" .pom) + files=() + types=() + classifiers=() + args=() + args+=(deploy:deploy-file) + args+=(-Durl=https://repository.apache.org/service/local/staging/deploy/maven2) + args+=(-DrepositoryId=apache.releases.https) + args+=(-DretryFailedDeploymentCount=10) + args+=(-DpomFile="${pom}") + if [ -f "${artifacts_dir}/${base}.jar" ]; then + jar="${artifacts_dir}/${base}.jar" + args+=(-Dfile="${jar}") + files+=("${signed_artifacts_dir}/${base}.jar.asc") + types+=("jar.asc") + classifiers+=("") + else + args+=(-Dfile="${pom}") + fi + files+=("${signed_artifacts_dir}/${base}.pom.asc") + types+=("pom.asc") + classifiers+=("") + if [ "$(echo "${artifacts_dir}/${base}"-*)" != "${artifacts_dir}/${base}-*" ]; then + for other_file in "${artifacts_dir}/${base}"-*; do + type="$(extract_type "${other_file}")" + case "${type}" in + sha256 | sha512) + continue + ;; + esac + classifier=$(extract_classifier "${other_file}" "${base}") + files+=("${other_file}") + types+=("${type}") + classifiers+=("${classifier}") + other_file_base="$(basename "${other_file}")" + files+=("${signed_artifacts_dir}/${other_file_base}.asc") + types+=("${type}.asc") + classifiers+=("${classifier}") + done + fi + args+=(-Dfiles="$( + IFS=, + echo "${files[*]}" + )") + args+=(-Dtypes="$( + IFS=, + echo "${types[*]}" + )") + args+=(-Dclassifiers="$( + IFS=, + echo "${classifiers[*]}" + )") + mvn "${args[@]}" + done + + echo + echo "Success!" + echo "Press the 'Close' button manually by Web interface:" + echo " https://repository.apache.org/#stagingRepositories" + echo "It publishes the artifacts to the staging repository:" + echo " https://repository.apache.org/content/repositories/staging/org/apache/arrow/" +fi + +echo +echo "Draft email for dev@arrow.apache.org mailing list" +echo "" +echo "---------------------------------------------------------" +cat < " + echo " e.g.: $0 19.0.1 1" + exit 1 +fi + +set -o pipefail +set -x + +VERSION="$1" +RC="$2" + +ARROW_DIST_BASE_URL="https://dist.apache.org/repos/dist/release/arrow" +DOWNLOAD_RC_BASE_URL="https://github.com/apache/arrow-java/releases/download/v${VERSION}-rc${RC}" +ARCHIVE_BASE_NAME="apache-arrow-java-${VERSION}" + +: "${VERIFY_DEFAULT:=1}" +: "${VERIFY_DOWNLOAD:=${VERIFY_DEFAULT}}" +: "${VERIFY_SIGN:=${VERIFY_DEFAULT}}" +: "${VERIFY_SOURCE:=${VERIFY_DEFAULT}}" +: "${VERIFY_BINARY:=${VERIFY_DEFAULT}}" + +VERIFY_SUCCESS=no + +setup_tmpdir() { + cleanup() { + if [ "${VERIFY_SUCCESS}" = "yes" ]; then + rm -rf "${VERIFY_TMPDIR}" + else + echo "Failed to verify release candidate. See ${VERIFY_TMPDIR} for details." + fi + } + + if [ -z "${VERIFY_TMPDIR:-}" ]; then + VERIFY_TMPDIR="$(mktemp -d -t "$1.XXXXX")" + trap cleanup EXIT + else + mkdir -p "${VERIFY_TMPDIR}" + fi +} + +download() { + curl \ + --fail \ + --location \ + --remote-name \ + --show-error \ + --silent \ + "$1" +} + +download_rc_file() { + if [ "${VERIFY_DOWNLOAD}" -gt 0 ]; then + download "${DOWNLOAD_RC_BASE_URL}/$1" + else + cp "${TOP_SOURCE_DIR}/$1" "$1" + fi +} + +import_gpg_keys() { + if [ "${VERIFY_SIGN}" -gt 0 ]; then + download "${ARROW_DIST_BASE_URL}/KEYS" + gpg --import KEYS + fi +} + +if type shasum >/dev/null 2>&1; then + sha256_verify="shasum -a 256 -c" + sha512_verify="shasum -a 512 -c" +else + sha256_verify="sha256sum -c" + sha512_verify="sha512sum -c" +fi + +fetch_archive() { + download_rc_file "${ARCHIVE_BASE_NAME}.tar.gz" + if [ "${VERIFY_SIGN}" -gt 0 ]; then + download_rc_file "${ARCHIVE_BASE_NAME}.tar.gz.asc" + gpg --verify "${ARCHIVE_BASE_NAME}.tar.gz.asc" "${ARCHIVE_BASE_NAME}.tar.gz" + fi + download_rc_file "${ARCHIVE_BASE_NAME}.tar.gz.sha256" + ${sha256_verify} "${ARCHIVE_BASE_NAME}.tar.gz.sha256" + download_rc_file "${ARCHIVE_BASE_NAME}.tar.gz.sha512" + ${sha512_verify} "${ARCHIVE_BASE_NAME}.tar.gz.sha512" +} + +ensure_source_directory() { + tar xf "${ARCHIVE_BASE_NAME}".tar.gz + + if [ -d "${TOP_SOURCE_DIR}/testing/data" ]; then + cp -a "${TOP_SOURCE_DIR}/testing" "${ARCHIVE_BASE_NAME}/" + else + git clone \ + https://github.com/apache/arrow-testing.git \ + "${ARCHIVE_BASE_NAME}/testing" + fi + + ARROW_TEST_DATA="$(pwd)/${ARCHIVE_BASE_NAME}/testing/data" + export ARROW_TEST_DATA +} + +test_source_distribution() { + if [ "${VERIFY_SOURCE}" -le 0 ]; then + return 0 + fi + + "${TOP_SOURCE_DIR}/ci/scripts/build.sh" "$(pwd)" build jni_dist + "${TOP_SOURCE_DIR}/ci/scripts/test.sh" "$(pwd)" build jni_dist + # TODO: JNI test + # TODO: Integration test +} + +test_binary_distribution() { + if [ "${VERIFY_BINARY}" -le 0 ]; then + return 0 + fi + + # TODO: jar test +} + +setup_tmpdir "arrow-java-${VERSION}-${RC}" +echo "Working in sandbox ${VERIFY_TMPDIR}" +cd "${VERIFY_TMPDIR}" + +import_gpg_keys +fetch_archive +ensure_source_directory +pushd "${ARCHIVE_BASE_NAME}" +test_source_distribution +test_binary_distribution +popd + +VERIFY_SUCCESS=yes +echo "RC looks good!" diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000000..a4de0bff18 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= -W +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000000..70c2ef2a37 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,28 @@ + + +# Documentation + +Build with Sphinx. + +```bash +cd docs +pip install -r requirements.txt +make html +``` diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000000..fa2d0bbe99 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,28 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +furo==2024.8.6 +myst-parser==4.0.0 +Sphinx==8.1.3 +sphinx-autobuild==2024.10.3 +sphinx-basic-ng==1.0.0b2 +sphinxcontrib-applehelp==2.0.0 +sphinxcontrib-devhelp==2.0.0 +sphinxcontrib-htmlhelp==2.1.0 +sphinxcontrib-jsmath==1.0.1 +sphinxcontrib-qthelp==2.0.0 +sphinxcontrib-serializinghtml==2.0.0 diff --git a/docs/source/_static/.gitignore b/docs/source/_static/.gitignore new file mode 100644 index 0000000000..e69de29bb2 diff --git a/docs/source/algorithm.rst b/docs/source/algorithm.rst new file mode 100644 index 0000000000..d4838967d6 --- /dev/null +++ b/docs/source/algorithm.rst @@ -0,0 +1,92 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +Java Algorithms +=============== + +Arrow's Java library provides algorithms for some commonly-used +functionalities. The algorithms are provided in the ``org.apache.arrow.algorithm`` +package of the ``algorithm`` module. + +Comparing Vector Elements +------------------------- + +Comparing vector elements is the basic for many algorithms. Vector +elements can be compared in one of the two ways: + +1. **Equality comparison**: there are two possible results for this type of comparisons: ``equal`` and ``unequal``. +Currently, this type of comparison is supported through the ``org.apache.arrow.vector.compare.VectorValueEqualizer`` +interface. + +2. **Ordering comparison**: there are three possible results for this type of comparisons: ``less than``, ``equal to`` +and ``greater than``. This comparison is supported by the abstract class ``org.apache.arrow.algorithm.sort.VectorValueComparator``. + +We provide default implementations to compare vector elements. However, users can also define ways +for customized comparisons. + +Vector Element Search +--------------------- + +A search algorithm tries to find a particular value in a vector. When successful, a vector index is +returned; otherwise, a ``-1`` is returned. The following search algorithms are provided: + +1. **Linear search**: this algorithm simply traverses the vector from the beginning, until a match is +found, or the end of the vector is reached. So it takes ``O(n)`` time, where ``n`` is the number of elements +in the vector. This algorithm is implemented in ``org.apache.arrow.algorithm.search.VectorSearcher#linearSearch``. + +2. **Binary search**: this represents a more efficient search algorithm, as it runs in ``O(log(n))`` time. +However, it is only applicable to sorted vectors. To get a sorted vector, +one can use one of our sorting algorithms, which will be discussed in the next section. This algorithm +is implemented in ``org.apache.arrow.algorithm.search.VectorSearcher#binarySearch``. + +3. **Parallel search**: when the vector is large, it takes a long time to traverse the elements to search +for a value. To make this process faster, one can split the vector into multiple partitions, and perform the +search for each partition in parallel. This is supported by ``org.apache.arrow.algorithm.search.ParallelSearcher``. + +4. **Range search**: for many scenarios, there can be multiple matching values in the vector. +If the vector is sorted, the matching values reside in a contiguous region in the vector. The +range search algorithm tries to find the upper/lower bound of the region in ``O(log(n))`` time. +An implementation is provided in ``org.apache.arrow.algorithm.search.VectorRangeSearcher``. + +Vector Sorting +-------------- + +Given a vector, a sorting algorithm turns it into a sorted one. The sorting criteria must +be specified by some ordering comparison operation. The sorting algorithms can be +classified into the following categories: + +1. **In-place sorter**: an in-place sorter performs the sorting by manipulating the original +vector, without creating any new vector. So it just returns the original vector after the sorting operations. +Currently, we have ``org.apache.arrow.algorithm.sort.FixedWidthInPlaceVectorSorter`` for in-place +sorting in ``O(nlog(n))`` time. As the name suggests, it only supports fixed width vectors. + +2. **Out-of-place sorter**: an out-of-place sorter does not mutate the original vector. Instead, +it copies vector elements to a new vector in sorted order, and returns the new vector. +We have ``org.apache.arrow.algorithm.sort.FixedWidthInPlaceVectorSorter.FixedWidthOutOfPlaceVectorSorter`` +and ``org.apache.arrow.algorithm.sort.FixedWidthInPlaceVectorSorter.VariableWidthOutOfPlaceVectorSorter`` +for fixed width and variable width vectors, respectively. Both algorithms run in ``O(nlog(n))`` time. + +3. **Index sorter**: this sorter does not actually sort the vector. Instead, it returns an integer +vector, which correspond to indices of vector elements in sorted order. With the index vector, one can +easily construct a sorted vector. In addition, some other tasks can be easily achieved, like finding the ``k`` th +smallest value in the vector. Index sorting is supported by ``org.apache.arrow.algorithm.sort.IndexSorter``, +which runs in ``O(nlog(n))`` time. It is applicable to vectors of any type. + +Other Algorithms +---------------- + +Other algorithms include vector deduplication, dictionary encoding, etc., in the ``algorithm`` module. diff --git a/docs/source/cdata.rst b/docs/source/cdata.rst new file mode 100644 index 0000000000..7b2924d259 --- /dev/null +++ b/docs/source/cdata.rst @@ -0,0 +1,468 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================ +C Data Interface +================ + +Arrow supports exchanging data without copying or serialization within the same process +through :external+arrow:ref:`c-data-interface`, even between different language runtimes. + +Java to Python +-------------- + +See :external+arrow:doc:`python/integration/python_java` to implement Java to +Python communication using the C Data Interface. + +Java to C++ +----------- + +See :external+arrow:doc:`developers/cpp/building` to build the Arrow C++ libraries: + +.. code-block:: shell + + $ git clone https://github.com/apache/arrow.git + $ cd arrow/cpp + $ mkdir build # from inside the `cpp` subdirectory + $ cd build + $ cmake .. --preset ninja-debug-minimal + $ cmake --build . + $ tree debug/ + debug/ + ├── libarrow.800.0.0.dylib + ├── libarrow.800.dylib -> libarrow.800.0.0.dylib + └── libarrow.dylib -> libarrow.800.dylib + +Share an Int64 array from C++ to Java +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**C++ Side** + +Implement a function in CDataCppBridge.h that exports an array via the C Data Interface: + +.. code-block:: cpp + + #include + #include + #include + + void FillInt64Array(const uintptr_t c_schema_ptr, const uintptr_t c_array_ptr) { + arrow::Int64Builder builder; + builder.Append(1); + builder.Append(2); + builder.Append(3); + builder.AppendNull(); + builder.Append(5); + builder.Append(6); + builder.Append(7); + builder.Append(8); + builder.Append(9); + builder.Append(10); + std::shared_ptr array = *builder.Finish(); + + struct ArrowSchema* c_schema = reinterpret_cast(c_schema_ptr); + auto c_schema_status = arrow::ExportType(*array->type(), c_schema); + if (!c_schema_status.ok()) c_schema_status.Abort(); + + struct ArrowArray* c_array = reinterpret_cast(c_array_ptr); + auto c_array_status = arrow::ExportArray(*array, c_array); + if (!c_array_status.ok()) c_array_status.Abort(); + } + +**Java Side** + +For this example, we will use `JavaCPP`_ to call our C++ function from Java, +without writing JNI bindings ourselves. + +.. code-block:: xml + + + + 4.0.0 + + org.example + java-cdata-example + 1.0-SNAPSHOT + + + 17 + 17 + 9.0.0 + + + + org.bytedeco + javacpp + 1.5.7 + + + org.apache.arrow + arrow-c-data + ${arrow.version} + + + org.apache.arrow + arrow-vector + ${arrow.version} + + + org.apache.arrow + arrow-memory-core + ${arrow.version} + + + org.apache.arrow + arrow-memory-netty + ${arrow.version} + + + org.apache.arrow + arrow-format + ${arrow.version} + + + + +.. code-block:: java + + import org.bytedeco.javacpp.annotation.Platform; + import org.bytedeco.javacpp.annotation.Properties; + import org.bytedeco.javacpp.tools.InfoMap; + import org.bytedeco.javacpp.tools.InfoMapper; + + @Properties( + target = "CDataJavaToCppExample", + value = @Platform( + include = { + "CDataCppBridge.h" + }, + compiler = {"cpp17"}, + linkpath = {"/arrow/cpp/build/debug/"}, + link = {"arrow"} + ) + ) + public class CDataJavaConfig implements InfoMapper { + + @Override + public void map(InfoMap infoMap) { + } + } + +.. code-block:: shell + + # Compile our Java code + $ javac -cp javacpp-1.5.7.jar CDataJavaConfig.java + + # Generate CDataInterfaceLibrary + $ java -jar javacpp-1.5.7.jar CDataJavaConfig.java + + # Generate libjniCDataInterfaceLibrary.dylib + $ java -jar javacpp-1.5.7.jar CDataJavaToCppExample.java + + # Validate libjniCDataInterfaceLibrary.dylib created + $ otool -L macosx-x86_64/libjniCDataJavaToCppExample.dylib + macosx-x86_64/libjniCDataJavaToCppExample.dylib: + libjniCDataJavaToCppExample.dylib (compatibility version 0.0.0, current version 0.0.0) + @rpath/libarrow.800.dylib (compatibility version 800.0.0, current version 800.0.0) + /usr/lib/libc++.1.dylib (compatibility version 1.0.0, current version 1200.3.0) + /usr/lib/libSystem.B.dylib (compatibility version 1.0.0, current version 1311.0.0) + +**Java Test** + +Let's create a Java class to test our bridge: + +.. code-block:: java + + import org.apache.arrow.c.ArrowArray; + import org.apache.arrow.c.ArrowSchema; + import org.apache.arrow.c.Data; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.BigIntVector; + + public class TestCDataInterface { + public static void main(String[] args) { + try( + BufferAllocator allocator = new RootAllocator(); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + ArrowArray arrowArray = ArrowArray.allocateNew(allocator) + ){ + CDataJavaToCppExample.FillInt64Array( + arrowSchema.memoryAddress(), arrowArray.memoryAddress()); + try( + BigIntVector bigIntVector = (BigIntVector) Data.importVector( + allocator, arrowArray, arrowSchema, null) + ){ + System.out.println("C++-allocated array: " + bigIntVector); + } + } + } + } + +.. code-block:: shell + + C++-allocated array: [1, 2, 3, null, 5, 6, 7, 8, 9, 10] + +Share an Int32 array from Java to C++ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Java Side** + +For this example, we will build a JAR with all dependencies bundled. + +.. code-block:: xml + + + + 4.0.0 + org.example + cpptojava + 1.0-SNAPSHOT + + 17 + 17 + 9.0.0 + + + + org.apache.arrow + arrow-c-data + ${arrow.version} + + + org.apache.arrow + arrow-memory-netty + ${arrow.version} + + + + + + org.apache.maven.plugins + maven-assembly-plugin + + + package + + single + + + + jar-with-dependencies + + + + + + + + + +.. code-block:: java + + import org.apache.arrow.c.ArrowArray; + import org.apache.arrow.c.ArrowSchema; + import org.apache.arrow.c.Data; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.FieldVector; + import org.apache.arrow.vector.IntVector; + import org.apache.arrow.vector.VectorSchemaRoot; + + import java.util.Arrays; + + public class ToBeCalledByCpp { + final static BufferAllocator allocator = new RootAllocator(); + + /** + * Create a {@link FieldVector} and export it via the C Data Interface + * @param schemaAddress Schema memory address to wrap + * @param arrayAddress Array memory address to wrap + */ + public static void fillVector(long schemaAddress, long arrayAddress){ + try (ArrowArray arrow_array = ArrowArray.wrap(arrayAddress); + ArrowSchema arrow_schema = ArrowSchema.wrap(schemaAddress) ) { + Data.exportVector(allocator, populateFieldVectorToExport(), null, arrow_array, arrow_schema); + } + } + + /** + * Create a {@link VectorSchemaRoot} and export it via the C Data Interface + * @param schemaAddress Schema memory address to wrap + * @param arrayAddress Array memory address to wrap + */ + public static void fillVectorSchemaRoot(long schemaAddress, long arrayAddress){ + try (ArrowArray arrow_array = ArrowArray.wrap(arrayAddress); + ArrowSchema arrow_schema = ArrowSchema.wrap(schemaAddress) ) { + Data.exportVectorSchemaRoot(allocator, populateVectorSchemaRootToExport(), null, arrow_array, arrow_schema); + } + } + + private static FieldVector populateFieldVectorToExport(){ + IntVector intVector = new IntVector("int-to-export", allocator); + intVector.allocateNew(3); + intVector.setSafe(0, 1); + intVector.setSafe(1, 2); + intVector.setSafe(2, 3); + intVector.setValueCount(3); + System.out.println("[Java] FieldVector: \n" + intVector); + return intVector; + } + + private static VectorSchemaRoot populateVectorSchemaRootToExport(){ + IntVector intVector = new IntVector("age-to-export", allocator); + intVector.setSafe(0, 10); + intVector.setSafe(1, 20); + intVector.setSafe(2, 30); + VectorSchemaRoot root = new VectorSchemaRoot(Arrays.asList(intVector)); + root.setRowCount(3); + System.out.println("[Java] VectorSchemaRoot: \n" + root.contentToTSVString()); + return root; + } + } + +Build the JAR and copy it to the C++ project. + +.. code-block:: shell + + $ mvn clean install + $ cp target/cpptojava-1.0-SNAPSHOT-jar-with-dependencies.jar /cpptojava.jar + +**C++ Side** + +This application uses JNI to call Java code, but transfers data (zero-copy) via the C Data Interface instead. + +.. code-block:: cpp + + #include + #include + + #include + #include + + JNIEnv *CreateVM(JavaVM **jvm) { + JNIEnv *env; + JavaVMInitArgs vm_args; + JavaVMOption options[2]; + options[0].optionString = "-Djava.class.path=cpptojava.jar"; + options[1].optionString = "-DXcheck:jni:pedantic"; + vm_args.version = JNI_VERSION_10; + vm_args.nOptions = 2; + vm_args.options = options; + int status = JNI_CreateJavaVM(jvm, (void **) &env, &vm_args); + if (status < 0) { + std::cerr << "\n<<<<< Unable to Launch JVM >>>>>\n" << std::endl; + return nullptr; + } + return env; + } + + int main() { + JNIEnv *env; + JavaVM *jvm; + env = CreateVM(&jvm); + if (env == nullptr) return EXIT_FAILURE; + jclass javaClassToBeCalledByCpp = env->FindClass("ToBeCalledByCpp"); + if (javaClassToBeCalledByCpp != nullptr) { + jmethodID fillVector = env->GetStaticMethodID(javaClassToBeCalledByCpp, + "fillVector", + "(JJ)V"); + if (fillVector != nullptr) { + struct ArrowSchema arrowSchema; + struct ArrowArray arrowArray; + std::cout << "\n<<<<< C++ to Java for Arrays >>>>>\n" << std::endl; + env->CallStaticVoidMethod(javaClassToBeCalledByCpp, fillVector, + static_cast(reinterpret_cast(&arrowSchema)), + static_cast(reinterpret_cast(&arrowArray))); + auto resultImportArray = arrow::ImportArray(&arrowArray, &arrowSchema); + std::shared_ptr array = resultImportArray.ValueOrDie(); + std::cout << "[C++] Array: " << array->ToString() << std::endl; + } else { + std::cerr << "Could not find fillVector method\n" << std::endl; + return EXIT_FAILURE; + } + jmethodID fillVectorSchemaRoot = env->GetStaticMethodID(javaClassToBeCalledByCpp, + "fillVectorSchemaRoot", + "(JJ)V"); + if (fillVectorSchemaRoot != nullptr) { + struct ArrowSchema arrowSchema; + struct ArrowArray arrowArray; + std::cout << "\n<<<<< C++ to Java for RecordBatch >>>>>\n" << std::endl; + env->CallStaticVoidMethod(javaClassToBeCalledByCpp, fillVectorSchemaRoot, + static_cast(reinterpret_cast(&arrowSchema)), + static_cast(reinterpret_cast(&arrowArray))); + auto resultImportVectorSchemaRoot = arrow::ImportRecordBatch(&arrowArray, &arrowSchema); + std::shared_ptr recordBatch = resultImportVectorSchemaRoot.ValueOrDie(); + std::cout << "[C++] RecordBatch: " << recordBatch->ToString() << std::endl; + } else { + std::cerr << "Could not find fillVectorSchemaRoot method\n" << std::endl; + return EXIT_FAILURE; + } + } else { + std::cout << "Could not find ToBeCalledByCpp class\n" << std::endl; + return EXIT_FAILURE; + } + jvm->DestroyJavaVM(); + return EXIT_SUCCESS; + } + +CMakeLists.txt definition file: + +.. code-block:: cmake + + cmake_minimum_required(VERSION 3.19) + project(cdatacpptojava) + find_package(JNI REQUIRED) + find_package(Arrow REQUIRED) + message(STATUS "Arrow version: ${ARROW_VERSION}") + include_directories(${JNI_INCLUDE_DIRS}) + set(CMAKE_CXX_STANDARD 17) + add_executable(${PROJECT_NAME} main.cpp) + target_link_libraries(cdatacpptojava PRIVATE Arrow::arrow_shared) + target_link_libraries(cdatacpptojava PRIVATE ${JNI_LIBRARIES}) + +**Result** + +.. code-block:: text + + <<<<< C++ to Java for Arrays >>>>> + [Java] FieldVector: + [1, 2, 3] + [C++] Array: [ + 1, + 2, + 3 + ] + + <<<<< C++ to Java for RecordBatch >>>>> + [Java] VectorSchemaRoot: + age-to-export + 10 + 20 + 30 + + [C++] RecordBatch: age-to-export: [ + 10, + 20, + 30 + ] + +.. _`JavaCPP`: https://github.com/bytedeco/javacpp diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 0000000000..3b17f326e1 --- /dev/null +++ b/docs/source/conf.py @@ -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. + +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +import pathlib +import xml.etree.ElementTree as ET + +project = 'arrow-java' +copyright = '2025, Apache Arrow Developers' +author = 'Apache Arrow Developers' + +top_level_pom_xml = pathlib.Path(__file__).parents[2] / "pom.xml" +tree = ET.parse(top_level_pom_xml) +ns = {"maven": "http://maven.apache.org/POM/4.0.0"} +release = tree.getroot().find("maven:version", ns).text + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = ["sphinx.ext.intersphinx"] + +templates_path = ['_templates'] +exclude_patterns = [] + +# -- Intersphinx ------------------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/extensions/intersphinx.html + +intersphinx_mapping = { + 'arrow': ('https://arrow.apache.org/docs/', None), + 'cookbook': ('https://arrow.apache.org/cookbook/java/', None), +} + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output + +html_theme = 'furo' +html_static_path = ['_static'] diff --git a/docs/source/dataset.rst b/docs/source/dataset.rst new file mode 100644 index 0000000000..deaa009576 --- /dev/null +++ b/docs/source/dataset.rst @@ -0,0 +1,309 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +======= +Dataset +======= + +.. warning:: + + Experimental: The Java module ``dataset`` is currently under early + development. API might be changed in each release of Apache Arrow until it + gets mature. + +Dataset is an universal layer in Apache Arrow for querying data in different +formats or in different partitioning strategies. Usually the data to be queried +is supposed to be located from a traditional file system, however Arrow Dataset +is not designed only for querying files but can be extended to serve all +possible data sources such as from inter-process communication or from other +network locations, etc. + +Getting Started +=============== + +Currently supported file formats are: + +- Apache Arrow (``.arrow``) +- Apache ORC (``.orc``) +- Apache Parquet (``.parquet``) +- Comma-Separated Values (``.csv``) +- Line-delimited JSON Values (``.json``) + +Below shows a simplest example of using Dataset to query a Parquet file in Java: + +.. code-block:: Java + + // read data from file /opt/example.parquet + String uri = "file:/opt/example.parquet"; + ScanOptions options = new ScanOptions(/*batchSize*/ 32768); + try ( + BufferAllocator allocator = new RootAllocator(); + DatasetFactory datasetFactory = new FileSystemDatasetFactory( + allocator, NativeMemoryPool.getDefault(), + FileFormat.PARQUET, uri); + Dataset dataset = datasetFactory.finish(); + Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches() + ) { + List batches = new ArrayList<>(); + while (reader.loadNextBatch()) { + try (VectorSchemaRoot root = reader.getVectorSchemaRoot()) { + final VectorUnloader unloader = new VectorUnloader(root); + batches.add(unloader.getRecordBatch()); + } + } + + // do something with read record batches, for example: + analyzeArrowData(batches); + + // finished the analysis of the data, close all resources: + AutoCloseables.close(batches); + } catch (Exception e) { + e.printStackTrace(); + } + +.. note:: + ``ArrowRecordBatch`` is a low-level composite Arrow data exchange format + that doesn't provide API to read typed data from it directly. + It's recommended to use utilities ``VectorLoader`` to load it into a schema + aware container ``VectorSchemaRoot`` by which user could be able to access + decoded data conveniently in Java. + + The ``ScanOptions batchSize`` argument takes effect only if it is set to a value + smaller than the number of rows in the recordbatch. + +.. seealso:: + Load record batches with :doc:`VectorSchemaRoot `. + +Schema +====== + +Schema of the data to be queried can be inspected via method +``DatasetFactory#inspect()`` before actually reading it. For example: + +.. code-block:: Java + + // read data from local file /opt/example.parquet + String uri = "file:/opt/example.parquet"; + BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + DatasetFactory factory = new FileSystemDatasetFactory(allocator, + NativeMemoryPool.getDefault(), FileFormat.PARQUET, uri); + + // inspect schema + Schema schema = factory.inspect(); + +For some of the data format that is compatible with a user-defined schema, user +can use method ``DatasetFactory#inspect(Schema schema)`` to create the dataset: + +.. code-block:: Java + + Schema schema = createUserSchema() + Dataset dataset = factory.finish(schema); + +Otherwise when the non-parameter method ``DatasetFactory#inspect()`` is called, +schema will be inferred automatically from data source. The same as the result +of ``DatasetFactory#inspect()``. + +Also, if projector is specified during scanning (see next section +:ref:`java-dataset-projection`), the actual schema of output data can be got +within method ``Scanner::schema()``: + +.. code-block:: Java + + Scanner scanner = dataset.newScan( + new ScanOptions(32768, Optional.of(new String[] {"id", "name"}))); + Schema projectedSchema = scanner.schema(); + +.. _java-dataset-projection: + +Projection (Subset of Columns) +============================== + +User can specify projections in ScanOptions. For example: + +.. code-block:: Java + + String[] projection = new String[] {"id", "name"}; + ScanOptions options = new ScanOptions(32768, Optional.of(projection)); + +If no projection is needed, leave the optional projection argument absent in +ScanOptions: + +.. code-block:: Java + + ScanOptions options = new ScanOptions(32768, Optional.empty()); + +Or use shortcut constructor: + +.. code-block:: Java + + ScanOptions options = new ScanOptions(32768); + +Then all columns will be emitted during scanning. + +Projection (Produce New Columns) and Filters +============================================ + +User can specify projections (new columns) or filters in ScanOptions using Substrait. For example: + +.. code-block:: Java + + ByteBuffer substraitExpressionFilter = getSubstraitExpressionFilter(); + ByteBuffer substraitExpressionProject = getSubstraitExpressionProjection(); + // Use Substrait APIs to create an Expression and serialize to a ByteBuffer + ScanOptions options = new ScanOptions.Builder(/*batchSize*/ 32768) + .columns(Optional.empty()) + .substraitExpressionFilter(substraitExpressionFilter) + .substraitExpressionProjection(getSubstraitExpressionProjection()) + .build(); + +.. seealso:: + + :doc:`Executing Projections and Filters Using Extended Expressions ` + Projections and Filters using Substrait. + +Read Data from HDFS +=================== + +``FileSystemDataset`` supports reading data from non-local file systems. HDFS +support is included in the official Apache Arrow Java package releases and +can be used directly without re-building the source code. + +To access HDFS data using Dataset API, pass a general HDFS URI to +``FilesSystemDatasetFactory``: + +.. code-block:: Java + + String uri = "hdfs://{hdfs_host}:{port}/data/example.parquet"; + BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + DatasetFactory factory = new FileSystemDatasetFactory(allocator, + NativeMemoryPool.getDefault(), FileFormat.PARQUET, uri); + +Native Memory Management +======================== + +To gain better performance and reduce code complexity, Java +``FileSystemDataset`` internally relies on C++ +``arrow::dataset::FileSystemDataset`` via JNI. +As a result, all Arrow data read from ``FileSystemDataset`` is supposed to be +allocated off the JVM heap. To manage this part of memory, an utility class +``NativeMemoryPool`` is provided to users. + +As a basic example, by using a listenable ``NativeMemoryPool``, user can pass +a listener hooking on C++ buffer allocation/deallocation: + +.. code-block:: Java + + AtomicLong reserved = new AtomicLong(0L); + ReservationListener listener = new ReservationListener() { + @Override + public void reserve(long size) { + reserved.getAndAdd(size); + } + + @Override + public void unreserve(long size) { + reserved.getAndAdd(-size); + } + }; + NativeMemoryPool pool = NativeMemoryPool.createListenable(listener); + FileSystemDatasetFactory factory = new FileSystemDatasetFactory(allocator, + pool, FileFormat.PARQUET, uri); + + +Also, it's a very common case to reserve the same amount of JVM direct memory +for the data read from datasets. For this use a built-in utility +class ``DirectReservationListener`` is provided: + +.. code-block:: Java + + NativeMemoryPool pool = NativeMemoryPool.createListenable( + DirectReservationListener.instance()); + +This way, once the allocated byte count of Arrow buffers reaches the limit of +JVM direct memory, ``OutOfMemoryError: Direct buffer memory`` will +be thrown during scanning. + +.. note:: + The default instance ``NativeMemoryPool.getDefaultMemoryPool()`` does + nothing on buffer allocation/deallocation. It's OK to use it in + the case of POC or testing, but for production use in complex environment, + it's recommended to manage memory by using a listenable memory pool. + +.. note:: + The ``BufferAllocator`` instance passed to ``FileSystemDatasetFactory``'s + constructor is also aware of the overall memory usage of the produced + dataset instances. Once the Java buffers are created the passed allocator + will become their parent allocator. + +Usage Notes +=========== + +Native Object Resource Management +--------------------------------- + +As another result of relying on JNI, all components related to +``FileSystemDataset`` should be closed manually or use try-with-resources to +release the corresponding native objects after using. For example: + +.. code-block:: Java + + String uri = "file:/opt/example.parquet"; + ScanOptions options = new ScanOptions(/*batchSize*/ 32768); + try ( + BufferAllocator allocator = new RootAllocator(); + DatasetFactory factory = new FileSystemDatasetFactory( + allocator, NativeMemoryPool.getDefault(), + FileFormat.PARQUET, uri); + Dataset dataset = factory.finish(); + Scanner scanner = dataset.newScan(options) + ) { + + // do something + + } catch (Exception e) { + e.printStackTrace(); + } + +If user forgets to close them then native object leakage might be caused. + +BatchSize +--------- + +The ``batchSize`` argument of ``ScanOptions`` is a limit on the size of an individual batch. + +For example, let's try to read a Parquet file with gzip compression and 3 row groups: + +.. code-block:: + + # Let configure ScanOptions as: + ScanOptions options = new ScanOptions(/*batchSize*/ 32768); + + $ parquet-tools meta data4_3rg_gzip.parquet + file schema: schema + age: OPTIONAL INT64 R:0 D:1 + name: OPTIONAL BINARY L:STRING R:0 D:1 + row group 1: RC:4 TS:182 OFFSET:4 + row group 2: RC:4 TS:190 OFFSET:420 + row group 3: RC:3 TS:179 OFFSET:838 + +Here, we set the batchSize in ScanOptions to 32768. Because that's greater +than the number of rows in the next batch, which is 4 rows because the first +row group has only 4 rows, then the program gets only 4 rows. The scanner +will not combine smaller batches to reach the limit, but it will split +large batches to stay under the limit. So in the case the row group had more +than 32768 rows, it would get split into blocks of 32768 rows or less. diff --git a/docs/source/developers/building.rst b/docs/source/developers/building.rst new file mode 100644 index 0000000000..b682957714 --- /dev/null +++ b/docs/source/developers/building.rst @@ -0,0 +1,624 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +.. highlight:: console + +.. _building-arrow-java: + +=================== +Building Arrow Java +=================== + +.. contents:: + +System Setup +============ + +Arrow Java uses the `Maven `_ build system. + +Building requires: + +* JDK 17+ +* Maven 3+ + +.. note:: + CI will test all supported JDK LTS versions, plus the latest non-LTS version. + +Building +======== + +All the instructions below assume that you have cloned the Arrow git +repository: + +.. code-block:: + + $ git clone https://github.com/apache/arrow.git + $ cd arrow + $ git submodule update --init --recursive + +These are the options available to compile Arrow Java modules with: + +* Maven build tool. +* Docker Compose. +* Archery. + +Building Java Modules +--------------------- + +To build the default modules, go to the project root and execute: + +Maven +~~~~~ + +.. code-block:: + + $ cd arrow/java + $ export JAVA_HOME= + $ java --version + $ mvn clean install + +Docker compose +~~~~~~~~~~~~~~ + +.. code-block:: + + $ cd arrow/java + $ export JAVA_HOME= + $ java --version + $ docker compose run java + +Archery +~~~~~~~ + +.. code-block:: + + $ cd arrow/java + $ export JAVA_HOME= + $ java --version + $ archery docker run java + +Building JNI Libraries (\*.dylib / \*.so / \*.dll) +-------------------------------------------------- + +First, we need to build the `C++ shared libraries`_ that the JNI bindings will use. +We can build these manually or we can use `Archery`_ to build them using a Docker container +(This will require installing Docker, Docker Compose, and Archery). + +.. note:: + If you are building on Apple Silicon, be sure to use a JDK version that was compiled + for that architecture. See, for example, the `Azul JDK `_. + + If you are building on Windows OS, see :ref:`Developing on Windows `. + +Maven +~~~~~ + +- To build only the JNI C Data Interface library (macOS / Linux): + + .. code-block:: text + + $ cd arrow/java + $ export JAVA_HOME= + $ java --version + $ mvn generate-resources -Pgenerate-libs-cdata-all-os -N + $ ls -latr ../java-dist/lib + |__ arrow_cdata_jni/ + +- To build only the JNI C Data Interface library (Windows): + + .. code-block:: + + $ cd arrow/java + $ mvn generate-resources -Pgenerate-libs-cdata-all-os -N + $ dir "../java-dist/bin" + |__ arrow_cdata_jni/ + +- To build all JNI libraries (macOS / Linux) except the JNI C Data Interface library: + + .. code-block:: text + + $ cd arrow/java + $ export JAVA_HOME= + $ java --version + $ mvn generate-resources -Pgenerate-libs-jni-macos-linux -N + $ ls -latr java-dist/lib + |__ arrow_dataset_jni/ + |__ arrow_orc_jni/ + |__ gandiva_jni/ + +- To build all JNI libraries (Windows) except the JNI C Data Interface library: + + .. code-block:: + + $ cd arrow/java + $ mvn generate-resources -Pgenerate-libs-jni-windows -N + $ dir "../java-dist/bin" + |__ arrow_dataset_jni/ + +CMake +~~~~~ + +- To build only the JNI C Data Interface library (macOS / Linux): + + .. code-block:: text + + $ cd arrow + $ mkdir -p java-dist java-cdata + $ cmake \ + -S java \ + -B java-cdata \ + -DARROW_JAVA_JNI_ENABLE_C=ON \ + -DARROW_JAVA_JNI_ENABLE_DEFAULT=OFF \ + -DBUILD_TESTING=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=java-dist + $ cmake --build java-cdata --target install --config Release + $ ls -latr java-dist/lib + |__ arrow_cdata_jni/ + +- To build only the JNI C Data Interface library (Windows): + + .. code-block:: text + + $ cd arrow + $ mkdir java-dist, java-cdata + $ cmake ^ + -S java ^ + -B java-cdata ^ + -DARROW_JAVA_JNI_ENABLE_C=ON ^ + -DARROW_JAVA_JNI_ENABLE_DEFAULT=OFF ^ + -DBUILD_TESTING=OFF ^ + -DCMAKE_BUILD_TYPE=Release ^ + -DCMAKE_INSTALL_PREFIX=java-dist + $ cmake --build java-cdata --target install --config Release + $ dir "java-dist/bin" + |__ arrow_cdata_jni/ + +- To build all JNI libraries (macOS / Linux) except the JNI C Data Interface library: + + .. code-block:: text + + $ cd arrow + $ brew bundle --file=cpp/Brewfile + # Homebrew Bundle complete! 25 Brewfile dependencies now installed. + $ brew uninstall aws-sdk-cpp + # (We can't use aws-sdk-cpp installed by Homebrew because it has + # an issue: https://github.com/aws/aws-sdk-cpp/issues/1809 ) + $ export JAVA_HOME= + $ mkdir -p java-dist cpp-jni + $ cmake \ + -S cpp \ + -B cpp-jni \ + -DARROW_BUILD_SHARED=OFF \ + -DARROW_CSV=ON \ + -DARROW_DATASET=ON \ + -DARROW_DEPENDENCY_SOURCE=BUNDLED \ + -DARROW_DEPENDENCY_USE_SHARED=OFF \ + -DARROW_FILESYSTEM=ON \ + -DARROW_GANDIVA=ON \ + -DARROW_GANDIVA_STATIC_LIBSTDCPP=ON \ + -DARROW_JSON=ON \ + -DARROW_ORC=ON \ + -DARROW_PARQUET=ON \ + -DARROW_S3=ON \ + -DARROW_SUBSTRAIT=ON \ + -DARROW_USE_CCACHE=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=java-dist \ + -DCMAKE_UNITY_BUILD=ON + $ cmake --build cpp-jni --target install --config Release + $ cmake \ + -S java \ + -B java-jni \ + -DARROW_JAVA_JNI_ENABLE_C=OFF \ + -DARROW_JAVA_JNI_ENABLE_DEFAULT=ON \ + -DBUILD_TESTING=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=java-dist \ + -DCMAKE_PREFIX_PATH=$PWD/java-dist \ + -DProtobuf_ROOT=$PWD/../cpp-jni/protobuf_ep-install \ + -DProtobuf_USE_STATIC_LIBS=ON + $ cmake --build java-jni --target install --config Release + $ ls -latr java-dist/lib/ + |__ arrow_dataset_jni/ + |__ arrow_orc_jni/ + |__ gandiva_jni/ + +- To build all JNI libraries (Windows) except the JNI C Data Interface library: + + .. code-block:: + + $ cd arrow + $ mkdir java-dist, cpp-jni + $ cmake ^ + -S cpp ^ + -B cpp-jni ^ + -DARROW_BUILD_SHARED=OFF ^ + -DARROW_CSV=ON ^ + -DARROW_DATASET=ON ^ + -DARROW_DEPENDENCY_USE_SHARED=OFF ^ + -DARROW_FILESYSTEM=ON ^ + -DARROW_GANDIVA=OFF ^ + -DARROW_JSON=ON ^ + -DARROW_ORC=ON ^ + -DARROW_PARQUET=ON ^ + -DARROW_S3=ON ^ + -DARROW_SUBSTRAIT=ON ^ + -DARROW_USE_CCACHE=ON ^ + -DARROW_WITH_BROTLI=ON ^ + -DARROW_WITH_LZ4=ON ^ + -DARROW_WITH_SNAPPY=ON ^ + -DARROW_WITH_ZLIB=ON ^ + -DARROW_WITH_ZSTD=ON ^ + -DCMAKE_BUILD_TYPE=Release ^ + -DCMAKE_INSTALL_PREFIX=java-dist ^ + -DCMAKE_UNITY_BUILD=ON ^ + -GNinja + $ cd cpp-jni + $ ninja install + $ cd ../ + $ cmake ^ + -S java ^ + -B java-jni ^ + -DARROW_JAVA_JNI_ENABLE_C=OFF ^ + -DARROW_JAVA_JNI_ENABLE_DATASET=ON ^ + -DARROW_JAVA_JNI_ENABLE_DEFAULT=ON ^ + -DARROW_JAVA_JNI_ENABLE_GANDIVA=OFF ^ + -DARROW_JAVA_JNI_ENABLE_ORC=ON ^ + -DBUILD_TESTING=OFF ^ + -DCMAKE_BUILD_TYPE=Release ^ + -DCMAKE_INSTALL_PREFIX=java-dist ^ + -DCMAKE_PREFIX_PATH=$PWD/java-dist + $ cmake --build java-jni --target install --config Release + $ dir "java-dist/bin" + |__ arrow_orc_jni/ + |__ arrow_dataset_jni/ + +Archery +~~~~~~~ + +.. code-block:: text + + $ cd arrow + $ archery docker run java-jni-manylinux-2014 + $ ls -latr java-dist + |__ arrow_cdata_jni/ + |__ arrow_dataset_jni/ + |__ arrow_orc_jni/ + |__ gandiva_jni/ + +Building Java JNI Modules +------------------------- + +- To compile the JNI bindings, use the ``arrow-c-data`` Maven profile: + + .. code-block:: + + $ cd arrow/java + $ mvn -Darrow.c.jni.dist.dir=/java-dist/lib -Parrow-c-data clean install + +- To compile the JNI bindings for ORC / Gandiva / Dataset, use the ``arrow-jni`` Maven profile: + + .. code-block:: + + $ cd arrow/java + $ mvn \ + -Darrow.cpp.build.dir=/java-dist/lib/ \ + -Darrow.c.jni.dist.dir=/java-dist/lib/ \ + -Parrow-jni clean install + +Testing +======= + +By default, Maven uses the same Java version to both build the code and run the tests. + +It is also possible to use a different JDK version for the tests. This requires Maven +toolchains to be configured beforehand, and then a specific test property needs to be set. + +Configuring Maven toolchains +---------------------------- + +To be able to use a JDK version for testing, it needs to be registered first in Maven ``toolchains.xml`` +configuration file usually located under ``${HOME}/.m2`` with the following snippet added to it: + + .. code-block:: + + + + + [...] + + + jdk + + 21 + temurin + + + path/to/jdk/home + + + + [...] + + + +Testing with a specific JDK +--------------------------- + +To run Arrow tests with a specific JDK version, use the ``arrow.test.jdk-version`` property. + +For example, to run Arrow tests with JDK 17, use the following snippet: + + .. code-block:: + + $ cd arrow/java + $ mvn -Darrow.test.jdk-version=17 clean verify + +IDE Configuration +================= + +IntelliJ +-------- + +To start working on Arrow in IntelliJ: build the project once from the command +line using ``mvn clean install``. Then open the ``java/`` subdirectory of the +Arrow repository, and update the following settings: + +* In the Files tool window, find the path ``vector/target/generated-sources``, + 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. +* 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 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"). +* If using IntelliJ's Maven integration to build, you may need to change + ```` to ``false`` in the pom.xml files due to an `IntelliJ bug + `__. +* 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 ``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 +IntelliJ Maven integration instead of with IntelliJ directly. + +Common Errors +============= + +* When working with the JNI code: if the C++ build cannot find dependencies, with errors like these: + + .. code-block:: + + Could NOT find Boost (missing: Boost_INCLUDE_DIR system filesystem) + Could NOT find Lz4 (missing: LZ4_LIB) + Could NOT find zstd (missing: ZSTD_LIB) + + Specify that the dependencies should be downloaded at build time (more details at `Dependency Resolution`_): + + .. code-block:: + + -Dre2_SOURCE=BUNDLED \ + -DBoost_SOURCE=BUNDLED \ + -Dutf8proc_SOURCE=BUNDLED \ + -DSnappy_SOURCE=BUNDLED \ + -DORC_SOURCE=BUNDLED \ + -DZLIB_SOURCE=BUNDLED + +.. _Archery: https://github.com/apache/arrow/blob/main/dev/archery/README.md +.. _Dependency Resolution: https://arrow.apache.org/docs/developers/cpp/building.html#individual-dependency-resolution +.. _C++ shared libraries: https://arrow.apache.org/docs/cpp/build_system.html + + +Installing Nightly Packages +=========================== + +.. warning:: + These packages are not official releases. Use them at your own risk. + +Arrow nightly builds are posted on the mailing list at `builds@arrow.apache.org`_. +The artifacts are uploaded to GitHub. For example, for 2022/07/30, they can be found at `GitHub Nightly`_. + + +Installing from Apache Nightlies +-------------------------------- +1. Look up the nightly version number for the Arrow libraries used. + + For example, for ``arrow-memory``, visit https://nightlies.apache.org/arrow/java/org/apache/arrow/arrow-memory/ and see what versions are available (e.g. 9.0.0.dev501). +2. Add Apache Nightlies Repository to the Maven/Gradle project. + + .. code-block:: xml + + + 9.0.0.dev501 + + ... + + + arrow-apache-nightlies + https://nightlies.apache.org/arrow/java + + + ... + + + org.apache.arrow + arrow-vector + ${arrow.version} + + + ... + +Installing Manually +------------------- + +1. Decide nightly packages repository to use, for example: https://github.com/ursacomputing/crossbow/releases/tag/nightly-packaging-2022-07-30-0-github-java-jars +2. Add packages to your pom.xml, for example: flight-core (it depends on: arrow-format, arrow-vector, arrow-memory-core and arrow-memory-netty). + + .. code-block:: xml + + + 17 + 17 + 9.0.0.dev501 + + + + + org.apache.arrow + flight-core + ${arrow.version} + + + +3. Download the necessary pom and jar files to a temporary directory: + + .. code-block:: shell + + $ mkdir nightly-packaging-2022-07-30-0-github-java-jars + $ cd nightly-packaging-2022-07-30-0-github-java-jars + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-java-root-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-format-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-format-9.0.0.dev501.jar + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-vector-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-vector-9.0.0.dev501.jar + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-memory-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-memory-core-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-memory-netty-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-memory-core-9.0.0.dev501.jar + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-memory-netty-9.0.0.dev501.jar + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/arrow-flight-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/flight-core-9.0.0.dev501.pom + $ wget https://github.com/ursacomputing/crossbow/releases/download/nightly-packaging-2022-07-30-0-github-java-jars/flight-core-9.0.0.dev501.jar + $ tree + . + ├── arrow-flight-9.0.0.dev501.pom + ├── arrow-format-9.0.0.dev501.jar + ├── arrow-format-9.0.0.dev501.pom + ├── arrow-java-root-9.0.0.dev501.pom + ├── arrow-memory-9.0.0.dev501.pom + ├── arrow-memory-core-9.0.0.dev501.jar + ├── arrow-memory-core-9.0.0.dev501.pom + ├── arrow-memory-netty-9.0.0.dev501.jar + ├── arrow-memory-netty-9.0.0.dev501.pom + ├── arrow-vector-9.0.0.dev501.jar + ├── arrow-vector-9.0.0.dev501.pom + ├── flight-core-9.0.0.dev501.jar + └── flight-core-9.0.0.dev501.pom + +4. Install the artifacts to the local Maven repository with ``mvn install:install-file``: + + .. code-block:: shell + + $ mvn install:install-file -Dfile="$(pwd)/arrow-java-root-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-java-root -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-format-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-format -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-format-9.0.0.dev501.jar" -DgroupId=org.apache.arrow -DartifactId=arrow-format -Dversion=9.0.0.dev501 -Dpackaging=jar + $ mvn install:install-file -Dfile="$(pwd)/arrow-vector-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-vector -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-vector-9.0.0.dev501.jar" -DgroupId=org.apache.arrow -DartifactId=arrow-vector -Dversion=9.0.0.dev501 -Dpackaging=jar + $ mvn install:install-file -Dfile="$(pwd)/arrow-memory-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-memory -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-memory-core-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-memory-core -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-memory-netty-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-memory-netty -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/arrow-memory-core-9.0.0.dev501.jar" -DgroupId=org.apache.arrow -DartifactId=arrow-memory-core -Dversion=9.0.0.dev501 -Dpackaging=jar + $ mvn install:install-file -Dfile="$(pwd)/arrow-memory-netty-9.0.0.dev501.jar" -DgroupId=org.apache.arrow -DartifactId=arrow-memory-netty -Dversion=9.0.0.dev501 -Dpackaging=jar + $ mvn install:install-file -Dfile="$(pwd)/arrow-flight-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=arrow-flight -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/flight-core-9.0.0.dev501.pom" -DgroupId=org.apache.arrow -DartifactId=flight-core -Dversion=9.0.0.dev501 -Dpackaging=pom + $ mvn install:install-file -Dfile="$(pwd)/flight-core-9.0.0.dev501.jar" -DgroupId=org.apache.arrow -DartifactId=flight-core -Dversion=9.0.0.dev501 -Dpackaging=jar + +5. Validate that the packages were installed: + + .. code-block:: shell + + $ tree ~/.m2/repository/org/apache/arrow + . + ├── arrow-flight + │   ├── 9.0.0.dev501 + │   │   └── arrow-flight-9.0.0.dev501.pom + ├── arrow-format + │   ├── 9.0.0.dev501 + │   │   ├── arrow-format-9.0.0.dev501.jar + │   │   └── arrow-format-9.0.0.dev501.pom + ├── arrow-java-root + │   ├── 9.0.0.dev501 + │   │   └── arrow-java-root-9.0.0.dev501.pom + ├── arrow-memory + │   ├── 9.0.0.dev501 + │   │   └── arrow-memory-9.0.0.dev501.pom + ├── arrow-memory-core + │   ├── 9.0.0.dev501 + │   │   ├── arrow-memory-core-9.0.0.dev501.jar + │   │   └── arrow-memory-core-9.0.0.dev501.pom + ├── arrow-memory-netty + │   ├── 9.0.0.dev501 + │   │   ├── arrow-memory-netty-9.0.0.dev501.jar + │   │   └── arrow-memory-netty-9.0.0.dev501.pom + ├── arrow-vector + │   ├── 9.0.0.dev501 + │   │   ├── _remote.repositories + │   │   ├── arrow-vector-9.0.0.dev501.jar + │   │   └── arrow-vector-9.0.0.dev501.pom + └── flight-core + ├── 9.0.0.dev501 + │   ├── flight-core-9.0.0.dev501.jar + │   └── flight-core-9.0.0.dev501.pom + +6. Compile your project like usual with ``mvn clean install``. + +.. _builds@arrow.apache.org: https://lists.apache.org/list.html?builds@arrow.apache.org +.. _GitHub Nightly: https://github.com/ursacomputing/crossbow/releases/tag/nightly-packaging-2022-07-30-0-github-java-jars + +Installing Staging Packages +=========================== + +.. warning:: + These packages are not official releases. Use them at your own risk. + +Arrow staging builds are created when a Release Candidate (RC) is being prepared. This allows users to test the RC in their applications before voting on the release. + + +Installing from Apache Staging +-------------------------------- +1. Look up the next version number for the Arrow libraries used. + +2. Add Apache Staging Repository to the Maven/Gradle project. + + .. code-block:: xml + + + 9.0.0 + + ... + + + arrow-apache-staging + https://repository.apache.org/content/repositories/staging + + + ... + + + org.apache.arrow + arrow-vector + ${arrow.version} + + + ... diff --git a/docs/source/developers/development.rst b/docs/source/developers/development.rst new file mode 100644 index 0000000000..dd1839257a --- /dev/null +++ b/docs/source/developers/development.rst @@ -0,0 +1,197 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +.. highlight:: console + +====================== +Development Guidelines +====================== + +.. contents:: + +Logger Abstraction +================== + +Apache Arrow Java uses the SLF4J API, so please configure SLF4J to see logs (e.g. via Logback/Apache Log4j): + +1. If no jar dependencies are added by the user via Logback or Apache Log4j then SLF4J will default + to no-operation (NOP) logging. + +2. If a user adds any dependencies via Logback or Apache Log4j but does not configure/add/define + logback.xml/log4j2.xml, then logs will default to DEBUG mode. + +3. To disable debug logs, the user must define their own rules within their logback.xml/log4j2.xml + and define their own loggers. + +Unit Testing +============ +Unit tests are run by Maven during the build. + +To speed up the build, you can skip them by passing -DskipTests. + +.. code-block:: + + $ cd arrow/java + $ mvn \ + -Darrow.cpp.build.dir=../java-dist/lib -Parrow-jni \ + -Darrow.c.jni.dist.dir=../java-dist/lib -Parrow-c-data \ + clean install + +Performance Testing +=================== + +The ``arrow-performance`` module contains benchmarks. + +Let's configure our environment to run performance tests: + +- Install `benchmark`_ +- Install `archery`_ + +In case you need to see your performance tests on the UI, then, configure (optional): + +- Install `conbench`_ + +Lets execute benchmark tests: + +.. code-block:: + + $ cd benchmarks + $ conbench java-micro --help + $ conbench java-micro + --iterations=1 + --commit=e90472e35b40f58b17d408438bb8de1641bfe6ef + --java-home= + --src= + --benchmark-filter=org.apache.arrow.adapter.AvroAdapterBenchmarks.testAvroToArrow + Benchmark Mode Cnt Score Error Units + AvroAdapterBenchmarks.testAvroToArrow avgt 725545.783 ns/op + Time to POST http://localhost:5000/api/login/ 0.14911699295043945 + Time to POST http://localhost:5000/api/benchmarks/ 0.06116318702697754 + +Then go to: http://127.0.0.1:5000/ to see reports: + +UI Home: + +.. image:: img/conbench_ui.png + +UI Runs: + +.. image:: img/conbench_runs.png + +UI Benchmark: + +.. image:: img/conbench_benchmark.png + +Integration Testing +=================== + +Integration tests can be run :ref:`via Archery `. +For example, assuming you only built Arrow Java and want to run the IPC +integration tests, you would do: + +.. code-block:: console + + $ archery integration --run-ipc --with-java 1 + +Code Style +========== + +The current Java code follows the `Google Java Style`_ with Apache license headers. + +Java code style is checked by `Spotless`_ during the build, and the continuous integration build will verify +that changes adhere to the style guide. + +Automatically fixing code style issues +-------------------------------------- + +- You can check the style without building the project with ``mvn spotless:check``. +- You can autoformat the source with ``mvn spotless:apply``. + +Example: + +.. code-block:: bash + + The following files had format violations: + src/main/java/org/apache/arrow/algorithm/rank/VectorRank.java + @@ -15,7 +15,6 @@ + ·*·limitations·under·the·License. + ·*/ + + - + package·org.apache.arrow.algorithm.rank; + + import·java.util.stream.IntStream; + Run 'mvn spotless:apply' to fix these violations. + +Code Formatter for Intellij IDEA and Eclipse +-------------------------------------------- + +Follow the instructions to set up google-java-format for: + +- `Eclipse`_ +- `IntelliJ`_ + + +Checkstyle +---------- + +Checkstyle is also used for general linting. The configuration is located at `checkstyle`_. +You can also just check the style without building the project. +This checks the code style of all source code under the current directory or from within an individual module. + +.. code-block:: + + $ mvn checkstyle:check + +Maven ``pom.xml`` style is enforced with Spotless using `Apache Maven pom.xml guidelines`_ +You can also just check the style without building the project. +This checks the style of all pom.xml files under the current directory or from within an individual module. + +.. code-block:: + + $ mvn spotless:check + +This applies the style to all pom.xml files under the current directory or from within an individual module. + +.. code-block:: + + $ mvn spotless:apply + +.. _benchmark: https://github.com/ursacomputing/benchmarks +.. _archery: https://github.com/apache/arrow/blob/main/dev/conbench_envs/README.md#L188 +.. _conbench: https://github.com/conbench/conbench +.. _checkstyle: https://github.com/apache/arrow/blob/main/java/dev/checkstyle/checkstyle.xml +.. _Apache Maven pom.xml guidelines: https://maven.apache.org/developers/conventions/code.html#pom-code-convention +.. _Spotless: https://github.com/diffplug/spotless +.. _Google Java Style: https://google.github.io/styleguide/javaguide.html +.. _Eclipse: https://github.com/google/google-java-format?tab=readme-ov-file#eclipse +.. _IntelliJ: https://github.com/google/google-java-format?tab=readme-ov-file#intellij-android-studio-and-other-jetbrains-ides + +Build Caching +============= + +Build caching is done through Develocity (formerly Maven Enterprise). To force +a build without the cache, run:: + + mvn clean install -Ddevelocity.cache.local.enabled=false -Ddevelocity.cache.remote.enabled=false + +This can be useful to make sure you see all warnings from ErrorProne, for example. + +ErrorProne +========== + +ErrorProne should be disabled for generated code. diff --git a/docs/source/developers/img/conbench_benchmark.png b/docs/source/developers/img/conbench_benchmark.png new file mode 100644 index 0000000000..3adf3e8cd4 Binary files /dev/null and b/docs/source/developers/img/conbench_benchmark.png differ diff --git a/docs/source/developers/img/conbench_runs.png b/docs/source/developers/img/conbench_runs.png new file mode 100644 index 0000000000..3a9c050776 Binary files /dev/null and b/docs/source/developers/img/conbench_runs.png differ diff --git a/docs/source/developers/img/conbench_ui.png b/docs/source/developers/img/conbench_ui.png new file mode 100644 index 0000000000..2f72df024b Binary files /dev/null and b/docs/source/developers/img/conbench_ui.png differ diff --git a/docs/source/developers/index.rst b/docs/source/developers/index.rst new file mode 100644 index 0000000000..976d1825cb --- /dev/null +++ b/docs/source/developers/index.rst @@ -0,0 +1,28 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +.. _java-development: + +**************** +Java Development +**************** + +.. toctree:: + :maxdepth: 2 + + building + development diff --git a/docs/source/flight.rst b/docs/source/flight.rst new file mode 100644 index 0000000000..fd0fdf07bc --- /dev/null +++ b/docs/source/flight.rst @@ -0,0 +1,239 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================ +Arrow Flight RPC +================ + +Arrow Flight is an RPC framework for efficient transfer of Arrow data +over the network. + +.. seealso:: + + :external+arrow:doc:`Flight protocol documentation ` + Documentation of the Flight protocol, including how to use + Flight conceptually. + + :external+cookbook:doc:`Java Cookbook ` + Recipes for using Arrow Flight in Java. + +Writing a Flight Service +======================== + +Flight servers implement the `FlightProducer`_ interface. For convenience, +they can subclass `NoOpFlightProducer`_ instead, which offers default +implementations of all the RPC methods. + +.. code-block:: Java + + public class TutorialFlightProducer implements FlightProducer { + @Override + // Override methods or use NoOpFlightProducer for only methods needed + } + +Each RPC method always takes a ``CallContext`` for common parameters. To indicate +failure, pass an exception to the "listener" if present, or else raise an +exception. + +.. code-block:: Java + + // Server + @Override + public void listFlights(CallContext context, Criteria criteria, StreamListener listener) { + // ... + listener.onError( + CallStatus.UNAUTHENTICATED.withDescription( + "Custom UNAUTHENTICATED description message.").toRuntimeException()); + // ... + } + + // Client + try{ + Iterable flightInfosBefore = flightClient.listFlights(Criteria.ALL); + // ... + } catch (FlightRuntimeException e){ + // Catch UNAUTHENTICATED exception + } + +To start a server, create a `Location`_ to specify where to listen, and then create +a `FlightServer`_ with an instance of a producer. This will start the server, but +won't block the rest of the program. Call ``FlightServer.awaitTermination`` +to block until the server stops. + +.. code-block:: Java + + class TutorialFlightProducer implements FlightProducer { + @Override + // Override methods or use NoOpFlightProducer for only methods needed + } + + Location location = Location.forGrpcInsecure("0.0.0.0", 0); + try( + BufferAllocator allocator = new RootAllocator(); + FlightServer server = FlightServer.builder( + allocator, + location, + new TutorialFlightProducer() + ).build(); + ){ + server.start(); + System.out.println("Server listening on port " + server.getPort()); + server.awaitTermination(); + } catch (Exception e) { + e.printStackTrace(); + } + +.. code-block:: shell + + Server listening on port 58104 + +Using the Flight Client +======================= + +To connect to a Flight service, create a `FlightClient`_ with a location. + +.. code-block:: Java + + Location location = Location.forGrpcInsecure("0.0.0.0", 58104); + + try(BufferAllocator allocator = new RootAllocator(); + FlightClient client = FlightClient.builder(allocator, location).build()){ + // ... Consume operations exposed by Flight server + } catch (Exception e) { + e.printStackTrace(); + } + +Cancellation and Timeouts +========================= + +When making a call, clients can optionally provide ``CallOptions``. This allows +clients to set a timeout on calls. Also, some objects returned by client RPC calls +expose a cancel method which allows terminating a call early. + +.. code-block:: Java + + Location location = Location.forGrpcInsecure("0.0.0.0", 58609); + + try(BufferAllocator allocator = new RootAllocator(); + FlightClient tutorialFlightClient = FlightClient.builder(allocator, location).build()){ + + Iterator resultIterator = tutorialFlightClient.doAction( + new Action("test-timeout"), + CallOptions.timeout(2, TimeUnit.SECONDS) + ); + } catch (Exception e) { + e.printStackTrace(); + } + +On the server side, timeouts are transparent. For cancellation, the server needs to manually poll +``setOnCancelHandler`` or ``isCancelled`` to check if the client has cancelled the call, +and if so, break out of any processing the server is currently doing. + +.. code-block:: Java + + // Client + Location location = Location.forGrpcInsecure("0.0.0.0", 58609); + try(BufferAllocator allocator = new RootAllocator(); + FlightClient tutorialFlightClient = FlightClient.builder(allocator, location).build()){ + try(FlightStream flightStream = flightClient.getStream(new Ticket(new byte[]{}))) { + // ... + flightStream.cancel("tutorial-cancel", new Exception("Testing cancellation option!")); + } + } catch (Exception e) { + e.printStackTrace(); + } + // Server + @Override + public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) { + // ... + listener.setOnCancelHandler(()->{ + // Implement logic to handle cancellation option + }); + } + +Enabling TLS +============ + +TLS can be enabled when setting up a server by providing a +certificate and key pair to ``FlightServer.Builder.useTls``. + +On the client side, use ``Location.forGrpcTls`` to create the Location for the client. + +Enabling Authentication +======================= + +.. warning:: Authentication is insecure without enabling TLS. + +Handshake-based authentication can be enabled by implementing +``ServerAuthHandler``. Authentication consists of two parts: on +initial client connection, the server and client authentication +implementations can perform any negotiation needed. The client authentication +handler then provides a token that will be attached to future calls. + +The client send data to be validated through ``ClientAuthHandler.authenticate`` +The server validate data received through ``ServerAuthHandler.authenticate``. + +Custom Middleware +================= + +Servers and clients support custom middleware (or interceptors) that are called on every +request and can modify the request in a limited fashion. These can be implemented by implementing the +``FlightServerMiddleware`` and ``FlightClientMiddleware`` interfaces. + +Middleware are fairly limited, but they can add headers to a +request/response. On the server, they can inspect incoming headers and +fail the request; hence, they can be used to implement custom +authentication methods. + +Adding Services +=============== + +Servers can add other gRPC services. For example, to add the `Health Check service `_: + +.. code-block:: Java + + final HealthStatusManager statusManager = new HealthStatusManager(); + final Consumer consumer = (builder) -> { + builder.addService(statusManager.getHealthService()); + }; + final Location location = forGrpcInsecure(LOCALHOST, 5555); + try ( + BufferAllocator a = new RootAllocator(Long.MAX_VALUE); + Producer producer = new Producer(a); + FlightServer s = FlightServer.builder(a, location, producer) + .transportHint("grpc.builderConsumer", consumer).build().start(); + ) { + Channel channel = NettyChannelBuilder.forAddress(location.toSocketAddress()).usePlaintext().build(); + HealthCheckResponse response = HealthGrpc + .newBlockingStub(channel) + .check(HealthCheckRequest.getDefaultInstance()); + + System.out.println(response.getStatus()); + } + + +:external+arrow:ref:`Flight best practices ` +=================================================================== + +See the :external+arrow:ref:`best practices for C++ `. + + +.. _`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 new file mode 100644 index 0000000000..09ce1dda0d --- /dev/null +++ b/docs/source/flight_sql.rst @@ -0,0 +1,32 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================ +Arrow Flight SQL +================ + +Arrow Flight SQL is an RPC framework for efficient transfer of Arrow data +over the network. + +.. seealso:: + + :external+arrow:doc:`Flight SQL protocol documentation ` + Documentation of the Flight SQL protocol. + +For usage information, see the `API documentation`_. + +.. _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 new file mode 100644 index 0000000000..6d40434a22 --- /dev/null +++ b/docs/source/flight_sql_jdbc_driver.rst @@ -0,0 +1,298 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +============================ +Arrow Flight SQL JDBC Driver +============================ + +The Flight SQL JDBC driver is a JDBC driver implementation that uses +the :external+arrow:doc:`Flight SQL protocol ` under +the hood. This driver can be used with any database that implements +Flight SQL. + +Installation and Requirements +============================= + +The driver is compatible with JDK 17+. Note that the following JVM +parameter is required: + +.. code-block:: shell + + java --add-opens=java.base/java.nio=ALL-UNNAMED ... + +To add a dependency via Maven, use a ``pom.xml`` like the following: + +.. code-block:: xml + + + + 4.0.0 + org.example + demo + 1.0-SNAPSHOT + + 18.1.0 + + + + org.apache.arrow + flight-sql-jdbc-driver + ${arrow.version} + + + + +Connecting to a Database +======================== + +The URI format is as follows:: + + jdbc:arrow-flight-sql://HOSTNAME:PORT[/?param1=val1¶m2=val2&...] + +For example, take this URI:: + + jdbc:arrow-flight-sql://localhost:12345/?username=admin&password=pass&useEncryption=1 + +This will connect to a Flight SQL service running on ``localhost`` on +port 12345. It will create a secure, encrypted connection, and +authenticate using the username ``admin`` and the password ``pass``. + +The components of the URI are as follows. + +* The URI scheme must be ``jdbc:arrow-flight-sql://``. +* **HOSTNAME** is the hostname of the Flight SQL service. +* **PORT** is the port of the Flight SQL service. + +Additional options can be passed as query parameters. Parameter names are +case-sensitive. The supported parameters are: + +.. list-table:: + :header-rows: 1 + + * - Parameter + - Default + - Description + + * - disableCertificateVerification + - false + - When TLS is enabled, whether to verify the server certificate + + * - password + - null + - The password for user/password authentication + + * - threadPoolSize + - 1 + - The size of an internal thread pool + + * - token + - null + - The token used for token authentication + + * - trustStore + - null + - When TLS is enabled, the path to the certificate store + + * - trustStorePassword + - null + - When TLS is enabled, the password for the certificate store + + * - tlsRootCerts + - null + - Path to PEM-encoded root certificates for TLS - use this as + an alternative to ``trustStore`` + + * - clientCertificate + - null + - Path to PEM-encoded client mTLS certificate when the Flight + SQL server requires client verification. + + * - clientKey + - null + - Path to PEM-encoded client mTLS key when the Flight + SQL server requires client verification. + + * - useEncryption + - true + - Whether to use TLS (the default is an encrypted connection) + + * - user + - null + - The username for user/password authentication + + * - useSystemTrustStore + - true + - When TLS is enabled, whether to use the system certificate store + + * - retainCookies + - true + - Whether to use cookies from the initial connection in subsequent + internal connections when retrieving streams from separate endpoints. + + * - retainAuth + - true + - Whether to use bearer tokens obtained from the initial connection + in subsequent internal connections used for retrieving streams + from separate endpoints. + +Note that URI values must be URI-encoded if they contain characters such +as !, @, $, etc. + +Any URI parameters that are not handled by the driver are passed to +the Flight SQL service as gRPC headers. For example, the following URI :: + + jdbc:arrow-flight-sql://localhost:12345/?useEncryption=0&database=mydb + +This will connect without authentication or encryption, to a Flight +SQL service running on ``localhost`` on port 12345. Each request will +also include a ``database=mydb`` gRPC header. + +Connection parameters may also be supplied using the Properties object +when using the JDBC Driver Manager to connect. When supplying using +the Properties object, values should *not* be URI-encoded. + +Parameters specified by the URI supercede parameters supplied by the +Properties object. When calling the `user/password overload of +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/index.rst b/docs/source/index.rst new file mode 100644 index 0000000000..5cdf41e197 --- /dev/null +++ b/docs/source/index.rst @@ -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. + +.. _java: + +Java Implementation +=================== + +This is the documentation of the Java API of Apache Arrow. For more details +on the Arrow format and other language bindings see the :doc:`parent documentation <../index>`. + +.. toctree:: + :maxdepth: 2 + + quickstartguide + overview + install + developers/index + + memory + vector + vector_schema_root + table + ipc + algorithm + flight + flight_sql + flight_sql_jdbc_driver + dataset + substrait + cdata + jdbc + Reference (javadoc) + Cookbook diff --git a/docs/source/install.rst b/docs/source/install.rst new file mode 100644 index 0000000000..e0b34515ef --- /dev/null +++ b/docs/source/install.rst @@ -0,0 +1,230 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +======================= +Installing Java Modules +======================= + +System Compatibility +==================== + +Java modules are regularly built and tested on macOS and Linux distributions. + +Java Compatibility +================== + +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: + +.. code-block:: shell + + # Directly on the command line + $ java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED -jar ... + # Indirectly via environment variables + $ env JDK_JAVA_OPTIONS="--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED" java -jar ... + +Otherwise, you may see errors like ``module java.base does not "opens +java.nio" to unnamed module`` or ``module java.base does not "opens +java.nio" to org.apache.arrow.memory.core`` + +Note that the command has changed from Arrow 15 and earlier. If you are still using the flags from that version +(``--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED``) you will see the +``module java.base does not "opens java.nio" to org.apache.arrow.memory.core`` error. + +If you are using flight-core or dependent modules, you will need to mark that flight-core can read unnamed modules. +Modifying the command above for Flight: + +.. code-block:: shell + + # Directly on the command line + $ java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED -jar ... + # Indirectly via environment variables + $ env JDK_JAVA_OPTIONS="--add-reads=org.apache.arrow.flight.core=ALL-UNNAMED --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED" java -jar ... + +Otherwise, you may see errors like ``java.lang.IllegalAccessError: superclass access check failed: class +org.apache.arrow.flight.ArrowMessage$ArrowBufRetainingCompositeByteBuf (in module org.apache.arrow.flight.core) +cannot access class io.netty.buffer.CompositeByteBuf (in unnamed module ...) because module +org.apache.arrow.flight.core does not read unnamed module ...`` + +Finally, if you are using arrow-dataset, you'll also need to report that JDK internals need to be exposed. +Modifying the command above for arrow-memory: + +.. code-block:: shell + + # Directly on the command line + $ java --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED -jar ... + # Indirectly via environment variables + $ env JDK_JAVA_OPTIONS="--add-opens=java.base/java.nio=org.apache.arrow.dataset,org.apache.arrow.memory.core,ALL-UNNAMED" java -jar ... + +Otherwise you may see errors such as ``java.lang.RuntimeException: java.lang.reflect.InaccessibleObjectException: +Unable to make static void java.nio.Bits.reserveMemory(long,long) accessible: module +java.base does not "opens java.nio" to module org.apache.arrow.dataset`` + +If using Maven and Surefire for unit testing, :ref:`this argument must +be added to Surefire as well `. + +Installing from Maven +===================== + +By default, Maven will download from the central repository: https://repo.maven.apache.org/maven2/org/apache/arrow/ + +Configure your pom.xml with the Java modules needed, for example: +arrow-vector, and arrow-memory-netty. + +.. code-block:: xml + + + + 4.0.0 + org.example + demo + 1.0-SNAPSHOT + + 9.0.0 + + + + org.apache.arrow + arrow-vector + ${arrow.version} + + + org.apache.arrow + arrow-memory-netty + ${arrow.version} + + + + +A bill of materials (BOM) module has been provided to simplify adding +Arrow modules. This eliminates the need to specify the version for +every module. An alternative to the above would be: + +.. code-block:: xml + + + + 4.0.0 + org.example + demo + 1.0-SNAPSHOT + + 15.0.0 + + + + org.apache.arrow + arrow-vector + + + org.apache.arrow + arrow-memory-netty + + + + + + org.apache.arrow + arrow-bom + ${arrow.version} + pom + import + + + + + +To use the Arrow Flight dependencies, also add the ``os-maven-plugin`` +plugin. This plugin generates useful platform-dependent properties +such as ``os.detected.name`` and ``os.detected.arch`` needed to resolve +transitive dependencies of Flight. + +.. code-block:: xml + + + + 4.0.0 + org.example + demo + 1.0-SNAPSHOT + + 9.0.0 + + + + org.apache.arrow + flight-core + ${arrow.version} + + + + + + kr.motd.maven + os-maven-plugin + 1.7.0 + + + + + +.. _java-install-maven-testing: + +The ``--add-opens`` flag must be added when running unit tests through Maven: + +.. code-block:: xml + + + + + org.apache.maven.plugins + maven-surefire-plugin + 3.0.0-M6 + + --add-opens=java.base/java.nio=ALL-UNNAMED + + + + + +Or they can be added via environment variable, for example when executing your code: + +.. code-block:: + + JDK_JAVA_OPTIONS="--add-opens=java.base/java.nio=ALL-UNNAMED" mvn exec:java -Dexec.mainClass="YourMainCode" + +Installing from Source +====================== + +See :ref:`java-development`. + +IDE Configuration +================= + +Generally, no additional configuration should be needed. However, +ensure your Maven or other build configuration has the ``--add-opens`` +flag as described above, so that the IDE picks it up and runs tests +with that flag as well. diff --git a/docs/source/ipc.rst b/docs/source/ipc.rst new file mode 100644 index 0000000000..f593917917 --- /dev/null +++ b/docs/source/ipc.rst @@ -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. + +=========================== +Reading/Writing IPC formats +=========================== +Arrow defines two types of binary formats for serializing record batches: + +* **Streaming format**: for sending an arbitrary number of record + batches. The format must be processed from start to end, and does not support + random access + +* **File or Random Access format**: for serializing a fixed number of record + batches. It supports random access, and thus is very useful when used with + memory maps + +Writing and Reading Streaming Format +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +First, let's populate a :class:`VectorSchemaRoot` with a small batch of records + +.. code-block:: Java + + BitVector bitVector = new BitVector("boolean", allocator); + VarCharVector varCharVector = new VarCharVector("varchar", allocator); + for (int i = 0; i < 10; i++) { + bitVector.setSafe(i, i % 2 == 0 ? 0 : 1); + varCharVector.setSafe(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); + } + bitVector.setValueCount(10); + varCharVector.setValueCount(10); + + List fields = Arrays.asList(bitVector.getField(), varCharVector.getField()); + List vectors = Arrays.asList(bitVector, varCharVector); + VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors); + +Now, we can begin writing a stream containing some number of these batches. For this we use :class:`ArrowStreamWriter` +(DictionaryProvider used for any vectors that are dictionary encoded is optional and can be null)) + +.. code-block:: Java + + try ( + ByteArrayOutputStream out = new ByteArrayOutputStream(); + ArrowStreamWriter writer = new ArrowStreamWriter(root, /*DictionaryProvider=*/null, Channels.newChannel(out)); + ) { + // ... do write into the ArrowStreamWriter + } + +Here we used an in-memory stream, but this could have been a socket or some other IO stream. Then we can do + +.. code-block:: Java + + writer.start(); + // write the first batch + writer.writeBatch(); + + // write another four batches. + for (int i = 0; i < 4; i++) { + // populate VectorSchemaRoot data and write the second batch + BitVector childVector1 = (BitVector)root.getVector(0); + VarCharVector childVector2 = (VarCharVector)root.getVector(1); + childVector1.reset(); + childVector2.reset(); + // ... do some populate work here, could be different for each batch + writer.writeBatch(); + } + + writer.end(); + +Note that, since the :class:`VectorSchemaRoot` in the writer is a container that can hold batches, batches flow through +:class:`VectorSchemaRoot` as part of a pipeline, so we need to populate data before ``writeBatch``, so that later batches +could overwrite previous ones. + +Now the :class:`ByteArrayOutputStream` contains the complete stream which contains 5 record batches. +We can read such a stream with :class:`ArrowStreamReader`. Note that the :class:`VectorSchemaRoot` within the reader +will be loaded with new values on every call to :class:`loadNextBatch()` + +.. code-block:: Java + + try (ArrowStreamReader reader = new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray()), allocator)) { + // This will be loaded with new values on every call to loadNextBatch + VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); + Schema schema = readRoot.getSchema(); + for (int i = 0; i < 5; i++) { + reader.loadNextBatch(); + // ... do something with readRoot + } + } + +Here we also give a simple example with dictionary encoded vectors + +.. code-block:: Java + + // create provider + DictionaryProvider.MapDictionaryProvider provider = new DictionaryProvider.MapDictionaryProvider(); + + try ( + final VarCharVector dictVector = new VarCharVector("dict", allocator); + final VarCharVector vector = new VarCharVector("vector", allocator); + ) { + // create dictionary vector + dictVector.allocateNewSafe(); + dictVector.setSafe(0, "aa".getBytes()); + dictVector.setSafe(1, "bb".getBytes()); + dictVector.setSafe(2, "cc".getBytes()); + dictVector.setValueCount(3); + + // create dictionary + Dictionary dictionary = + new Dictionary(dictVector, new DictionaryEncoding(1L, false, /*indexType=*/null)); + provider.put(dictionary); + + // create original data vector + vector.allocateNewSafe(); + vector.setSafe(0, "bb".getBytes()); + vector.setSafe(1, "bb".getBytes()); + vector.setSafe(2, "cc".getBytes()); + vector.setSafe(3, "aa".getBytes()); + vector.setValueCount(4); + + // get the encoded vector + IntVector encodedVector = (IntVector) DictionaryEncoder.encode(vector, dictionary); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + + // create VectorSchemaRoot + List fields = Arrays.asList(encodedVector.getField()); + List vectors = Arrays.asList(encodedVector); + try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors)) { + + // write data + ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)); + writer.start(); + writer.writeBatch(); + writer.end(); + } + + // read data + try (ArrowStreamReader reader = new ArrowStreamReader(new ByteArrayInputStream(out.toByteArray()), allocator)) { + reader.loadNextBatch(); + VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); + // get the encoded vector + IntVector intVector = (IntVector) readRoot.getVector(0); + + // get dictionaries and decode the vector + Map dictionaryMap = reader.getDictionaryVectors(); + long dictionaryId = intVector.getField().getDictionary().getId(); + try (VarCharVector varCharVector = + (VarCharVector) DictionaryEncoder.decode(intVector, dictionaryMap.get(dictionaryId))) { + // ... use decoded vector + } + } + } + +Writing and Reading Random Access Files +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The :class:`ArrowFileWriter` has the same API as :class:`ArrowStreamWriter` + +.. code-block:: Java + + try ( + ByteArrayOutputStream out = new ByteArrayOutputStream(); + ArrowFileWriter writer = new ArrowFileWriter(root, /*DictionaryProvider=*/null, Channels.newChannel(out)); + ) { + writer.start(); + // write the first batch + writer.writeBatch(); + // write another four batches. + for (int i = 0; i < 4; i++) { + // ... do populate work + writer.writeBatch(); + } + writer.end(); + } + +The difference between :class:`ArrowFileReader` and :class:`ArrowStreamReader` is that the input source +must have a ``seek`` method for random access. Because we have access to the entire payload, we know the +number of record batches in the file, and can read any at random + +.. code-block:: Java + + try (ArrowFileReader reader = new ArrowFileReader( + new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) { + + // read the 4-th batch + ArrowBlock block = reader.getRecordBlocks().get(3); + reader.loadRecordBatch(block); + VectorSchemaRoot readBatch = reader.getVectorSchemaRoot(); + } diff --git a/docs/source/jdbc.rst b/docs/source/jdbc.rst new file mode 100644 index 0000000000..e054127faa --- /dev/null +++ b/docs/source/jdbc.rst @@ -0,0 +1,284 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================== +Arrow JDBC Adapter +================== + +The Arrow JDBC Adapter assists with working with JDBC and Arrow +data. Currently, it supports reading JDBC ResultSets into Arrow +VectorSchemaRoots. + +ResultSet to VectorSchemaRoot Conversion +======================================== + +This can be accessed via the JdbcToArrow class. The resulting +ArrowVectorIterator will convert a ResultSet to Arrow data in batches +of rows. + +.. code-block:: java + + try (ArrowVectorIterator it = JdbcToArrow.sqlToArrowVectorIterator(resultSet, allocator)) { + while (it.hasNext()) { + VectorSchemaRoot root = it.next(); + // Consume the root… + } + } + +The batch size and type mapping can both be customized: + +.. code-block:: java + + JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, /*calendar=*/null) + .setReuseVectorSchemaRoot(reuseVectorSchemaRoot) + .setJdbcToArrowTypeConverter((jdbcFieldInfo -> { + switch (jdbcFieldInfo.getJdbcType()) { + case Types.BIGINT: + // Assume actual value range is SMALLINT + return new ArrowType.Int(16, true); + default: + return null; + } + })) + .build(); + try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { + while (iter.hasNext()) { + VectorSchemaRoot root = iter.next(); + // Consume the root… + } + } + +The JDBC type can be explicitly specified, which is useful since JDBC +drivers can give spurious type information. For example, the Postgres +driver has been observed to use Decimal types with scale and precision +0; these cases can be handled by specifying the type explicitly before +reading. Also, some JDBC drivers may return BigDecimal values with +inconsistent scale. A RoundingMode can be set to handle these cases: + +.. code-block:: java + + Map mapping = new HashMap<>(); + mapping.put(1, new JdbcFieldInfo(Types.DECIMAL, 20, 7)); + JdbcToArrowConfig config = new JdbcToArrowConfigBuilder(allocator, /*calendar=*/null) + .setBigDecimalRoundingMode(RoundingMode.UNNECESSARY) + .setExplicitTypesByColumnIndex(mapping) + .build(); + try (ArrowVectorIterator iter = JdbcToArrow.sqlToArrowVectorIterator(rs, config)) { + while (iter.hasNext()) { + VectorSchemaRoot root = iter.next(); + // Consume the root… + } + } + +The mapping from JDBC type to Arrow type can be overridden via the +``JdbcToArrowConfig``, but it is not possible to customize the +conversion from JDBC value to Arrow value itself, nor is it possible +to define a conversion for an unsupported type. + +Type Mapping +------------ + +The JDBC to Arrow type mapping can be obtained at runtime from +`JdbcToArrowUtils.getArrowTypeFromJdbcType`_. + +.. _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 | ++====================+====================+=======+ +| ARRAY | List | \(1) | ++--------------------+--------------------+-------+ +| BIGINT | Int64 | | ++--------------------+--------------------+-------+ +| BINARY | Binary | | ++--------------------+--------------------+-------+ +| BIT | Bool | | ++--------------------+--------------------+-------+ +| BLOB | Binary | | ++--------------------+--------------------+-------+ +| BOOLEAN | Bool | | ++--------------------+--------------------+-------+ +| CHAR | Utf8 | | ++--------------------+--------------------+-------+ +| CLOB | Utf8 | | ++--------------------+--------------------+-------+ +| DATE | Date32 | | ++--------------------+--------------------+-------+ +| DECIMAL | Decimal128 | \(2) | ++--------------------+--------------------+-------+ +| DOUBLE | Double | | ++--------------------+--------------------+-------+ +| FLOAT | Float32 | | ++--------------------+--------------------+-------+ +| INTEGER | Int32 | | ++--------------------+--------------------+-------+ +| LONGVARBINARY | Binary | | ++--------------------+--------------------+-------+ +| LONGNVARCHAR | Utf8 | | ++--------------------+--------------------+-------+ +| LONGVARCHAR | Utf8 | | ++--------------------+--------------------+-------+ +| NCHAR | Utf8 | | ++--------------------+--------------------+-------+ +| NULL | Null | | ++--------------------+--------------------+-------+ +| NUMERIC | Decimal128 | | ++--------------------+--------------------+-------+ +| NVARCHAR | Utf8 | | ++--------------------+--------------------+-------+ +| REAL | Float32 | | ++--------------------+--------------------+-------+ +| SMALLINT | Int16 | | ++--------------------+--------------------+-------+ +| STRUCT | Struct | \(3) | ++--------------------+--------------------+-------+ +| TIME | Time32[ms] | | ++--------------------+--------------------+-------+ +| TIMESTAMP | Timestamp[ms] | \(4) | ++--------------------+--------------------+-------+ +| TINYINT | Int8 | | ++--------------------+--------------------+-------+ +| VARBINARY | Binary | | ++--------------------+--------------------+-------+ +| VARCHAR | Utf8 | | ++--------------------+--------------------+-------+ + +* \(1) The list value type must be explicitly configured and cannot be + inferred. Use `setArraySubTypeByColumnIndexMap`_ or + `setArraySubTypeByColumnNameMap`_. +* \(2) By default, the scale of decimal values must match the scale in + the type exactly; precision is allowed to be any value greater or + equal to the type precision. If there is a mismatch, by default, an + exception will be thrown. This can be configured by setting a + different RoundingMode with setBigDecimalRoundingMode. +* \(3) Not fully supported: while the type conversion is defined, the + value conversion is not. See ARROW-17006_. +* \(4) If a Calendar is provided, then the timestamp will have the + timezone of the calendar, else it will be a timestamp without + timezone. + +.. _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 +========================================================== + +The adapter can bind rows of Arrow data from a VectorSchemaRoot to +parameters of a JDBC PreparedStatement. This can be accessed via the +JdbcParameterBinder class. Each call to next() will bind parameters +from the next row of data, and then the application can execute the +statement, call addBatch(), etc. as desired. Null values will lead to +a setNull call with an appropriate JDBC type code (listed below). + +.. code-block:: java + + final JdbcParameterBinder binder = + JdbcParameterBinder.builder(statement, root).bindAll().build(); + while (binder.next()) { + statement.executeUpdate(); + } + // Use a VectorLoader to update the root + binder.reset(); + while (binder.next()) { + statement.executeUpdate(); + } + +The mapping of vectors to parameters, the JDBC type code used by the +converters, and the type conversions themselves can all be customized: + +.. code-block:: java + + final JdbcParameterBinder binder = + JdbcParameterBinder.builder(statement, root) + .bind(/*parameterIndex*/2, /*columnIndex*/0) + .bind(/*parameterIndex*/1, customColumnBinderInstance) + .build(); + +Type Mapping +------------ + +The Arrow to JDBC type mapping can be obtained at runtime via +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 | ++============================+============================+=======+ +| Binary | VARBINARY (setBytes) | | ++----------------------------+----------------------------+-------+ +| Bool | BOOLEAN (setBoolean) | | ++----------------------------+----------------------------+-------+ +| Date32 | DATE (setDate) | | ++----------------------------+----------------------------+-------+ +| Date64 | DATE (setDate) | | ++----------------------------+----------------------------+-------+ +| Decimal128 | DECIMAL (setBigDecimal) | | ++----------------------------+----------------------------+-------+ +| Decimal256 | DECIMAL (setBigDecimal) | | ++----------------------------+----------------------------+-------+ +| FixedSizeBinary | BINARY (setBytes) | | ++----------------------------+----------------------------+-------+ +| Uuid (extension) | OTHER (setObject) | \(3) | ++----------------------------+----------------------------+-------+ +| Float32 | REAL (setFloat) | | ++----------------------------+----------------------------+-------+ +| Int8 | TINYINT (setByte) | | ++----------------------------+----------------------------+-------+ +| Int16 | SMALLINT (setShort) | | ++----------------------------+----------------------------+-------+ +| Int32 | INTEGER (setInt) | | ++----------------------------+----------------------------+-------+ +| Int64 | BIGINT (setLong) | | ++----------------------------+----------------------------+-------+ +| LargeBinary | LONGVARBINARY (setBytes) | | ++----------------------------+----------------------------+-------+ +| LargeUtf8 | LONGVARCHAR (setString) | \(1) | ++----------------------------+----------------------------+-------+ +| Time[s] | TIME (setTime) | | ++----------------------------+----------------------------+-------+ +| Time[ms] | TIME (setTime) | | ++----------------------------+----------------------------+-------+ +| Time[us] | TIME (setTime) | | ++----------------------------+----------------------------+-------+ +| Time[ns] | TIME (setTime) | | ++----------------------------+----------------------------+-------+ +| Timestamp[s] | TIMESTAMP (setTimestamp) | \(2) | ++----------------------------+----------------------------+-------+ +| Timestamp[ms] | TIMESTAMP (setTimestamp) | \(2) | ++----------------------------+----------------------------+-------+ +| Timestamp[us] | TIMESTAMP (setTimestamp) | \(2) | ++----------------------------+----------------------------+-------+ +| Timestamp[ns] | TIMESTAMP (setTimestamp) | \(2) | ++----------------------------+----------------------------+-------+ +| Utf8 | VARCHAR (setString) | | ++----------------------------+----------------------------+-------+ + +* \(1) Strings longer than Integer.MAX_VALUE bytes (the maximum length + of a Java ``byte[]``) will cause a runtime exception. +* \(2) If the timestamp has a timezone, the JDBC type defaults to + TIMESTAMP_WITH_TIMEZONE. If the timestamp has no timezone, + technically there is not a correct conversion from Arrow value to + 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 new file mode 100644 index 0000000000..5b9148223a --- /dev/null +++ b/docs/source/memory.rst @@ -0,0 +1,499 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================= +Memory Management +================= + +The memory modules contain all the functionality that Arrow uses to allocate and deallocate memory. This document is divided in two parts: +The first part, *Memory Basics*, provides a high-level introduction. The following section, *Arrow Memory In-Depth*, fills in the details. + +Memory Basics +============= +This section will introduce you to the major concepts in Java’s memory management: + +* `ArrowBuf`_ +* `BufferAllocator`_ +* Reference counting + +It also provides some guidelines for working with memory in Arrow, and describes how to debug memory issues when they arise. + +Getting Started +--------------- + +Arrow's memory management is built around the needs of the columnar format and using off-heap memory. +Arrow Java has its own independent implementation. It does not wrap the C++ implementation, although the framework is flexible enough +to be used with memory allocated in C++ that is used by Java code. + +Arrow provides multiple modules: the core interfaces, and implementations of the interfaces. +Users need the core interfaces, and exactly one of the implementations. + +* ``memory-core``: Provides the interfaces used by the Arrow libraries and applications. +* ``memory-netty``: An implementation of the memory interfaces based on the `Netty`_ library. +* ``memory-unsafe``: An implementation of the memory interfaces based on the `sun.misc.Unsafe`_ library. + + +ArrowBuf +-------- + +ArrowBuf represents a single, contiguous region of `direct memory`_. It consists of an address and a length, +and provides low-level interfaces for working with the contents, similar to ByteBuffer. + +Unlike (Direct)ByteBuffer, it has reference counting built in, as discussed later. + +Why Arrow Uses Direct Memory +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +* The JVM can optimize I/O operations when using direct memory/direct buffers; it will attempt to avoid copying buffer contents to/from an intermediate buffer. This can speed up IPC in Arrow. +* Since Arrow always uses direct memory, JNI modules can directly wrap native memory addresses instead of copying data. We use this in modules like the C Data Interface. +* Conversely, on the C++ side of the JNI boundary, we can directly access the memory in ArrowBuf without copying data. + +BufferAllocator +--------------- + +The `BufferAllocator`_ is primarily an arena or nursery used for accounting of buffers (ArrowBuf instances). +As the name suggests, it can allocate new buffers associated with itself, but it can also +handle the accounting for buffers allocated elsewhere. For example, it handles the Java-side accounting for +memory allocated in C++ and shared with Java using the C-Data Interface. In the code below it performs an allocation: + +.. code-block:: Java + + import org.apache.arrow.memory.ArrowBuf; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + + try(BufferAllocator bufferAllocator = new RootAllocator(8 * 1024)){ + ArrowBuf arrowBuf = bufferAllocator.buffer(4 * 1024); + System.out.println(arrowBuf); + arrowBuf.close(); + } + +.. code-block:: shell + + ArrowBuf[2], address:140363641651200, length:4096 + +The concrete implementation of the BufferAllocator interface is `RootAllocator`_. Applications should generally create +one RootAllocator at the start of the program, and use it through the BufferAllocator interface. Allocators implement +AutoCloseable and must be closed after the application is done with them; this will check that all outstanding memory +has been freed (see the next section). + +Arrow provides a tree-based model for memory allocation. The RootAllocator is created first, then more allocators +are created as children of an existing allocator via `newChildAllocator`_. When creating a RootAllocator or a child +allocator, a memory limit is provided, and when allocating memory, the limit is checked. Furthermore, when allocating +memory from a child allocator, those allocations are also reflected in all parent allocators. Hence, the RootAllocator +effectively sets the program-wide memory limit, and serves as the master bookkeeper for all memory allocations. + +Child allocators are not strictly required, but can help better organize code. For instance, a lower memory limit can +be set for a particular section of code. The child allocator can be closed when that section completes, +at which point it checks that that section didn't leak any memory. +Child allocators can also be named, which makes it easier to tell where an ArrowBuf came from during debugging. + +Reference counting +------------------ + +Because direct memory is expensive to allocate and deallocate, allocators may share direct buffers. To manage shared buffers +deterministically, we use manual reference counting instead of the garbage collector. +This simply means that each buffer has a counter keeping track of the number of references to +the buffer, and the user is responsible for properly incrementing/decrementing the counter as the buffer is used. + +In Arrow, each ArrowBuf has an associated `ReferenceManager`_ that tracks the reference count. You can retrieve +it with ArrowBuf.getReferenceManager(). The reference count is updated using `ReferenceManager.release`_ to decrement the count, +and `ReferenceManager.retain`_ to increment it. + +Of course, this is tedious and error-prone, so instead of directly working with buffers, we typically use +higher-level APIs like ValueVector. Such classes generally implement Closeable/AutoCloseable and will automatically +decrement the reference count when closed. + +Allocators implement AutoCloseable as well. In this case, closing the allocator will check that all buffers +obtained from the allocator are closed. If not, ``close()`` method will raise an exception; this helps track +memory leaks from unclosed buffers. + +Reference counting needs to be handled carefully. To ensure that an +independent section of code has fully cleaned up all allocated buffers, use a new child allocator. + +Development Guidelines +---------------------- + +Applications should generally: + +* Use the BufferAllocator interface in APIs instead of RootAllocator. +* Create one RootAllocator at the start of the program and explicitly pass it when needed. +* ``close()`` allocators after use (whether they are child allocators or the RootAllocator), either manually or preferably via a try-with-resources statement. + + +Debugging Memory Leaks/Allocation +--------------------------------- + +In ``DEBUG`` mode, the allocator and supporting classes will record additional +debug tracking information to better track down memory leaks and issues. To +enable DEBUG mode pass the following system property to the VM when starting +``-Darrow.memory.debug.allocator=true``. + +When DEBUG is enabled, a log will be kept of allocations. Configure SLF4J to see these logs (e.g. via Logback/Apache Log4j). +Consider the following example to see how it helps us with the tracking of allocators: + +.. code-block:: Java + + import org.apache.arrow.memory.ArrowBuf; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + + try (BufferAllocator bufferAllocator = new RootAllocator(8 * 1024)) { + ArrowBuf arrowBuf = bufferAllocator.buffer(4 * 1024); + System.out.println(arrowBuf); + } + +Without the debug mode enabled, when we close the allocator, we get this: + +.. code-block:: shell + + 11:56:48.944 [main] INFO o.apache.arrow.memory.BaseAllocator - Debug mode disabled. + ArrowBuf[2], address:140508391276544, length:4096 + 16:28:08.847 [main] ERROR o.apache.arrow.memory.BaseAllocator - Memory was leaked by query. Memory leaked: (4096) + Allocator(ROOT) 0/4096/4096/8192 (res/actual/peak/limit) + +Enabling the debug mode, we get more details: + +.. code-block:: shell + + 11:56:48.944 [main] INFO o.apache.arrow.memory.BaseAllocator - Debug mode enabled. + ArrowBuf[2], address:140437894463488, length:4096 + Exception in thread "main" java.lang.IllegalStateException: Allocator[ROOT] closed with outstanding buffers allocated (1). + Allocator(ROOT) 0/4096/4096/8192 (res/actual/peak/limit) + child allocators: 0 + ledgers: 1 + ledger[1] allocator: ROOT), isOwning: , size: , references: 1, life: 261438177096661..0, allocatorManager: [, life: ] holds 1 buffers. + ArrowBuf[2], address:140437894463488, length:4096 + reservations: 0 + +Additionally, in debug mode, `ArrowBuf.print()`_ can be used to obtain a debug string. +This will include information about allocation operations on the buffer with stack traces, such as when/where the buffer was allocated. + +.. code-block:: java + + import org.apache.arrow.memory.ArrowBuf; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + + try (final BufferAllocator allocator = new RootAllocator()) { + try (final ArrowBuf buf = allocator.buffer(1024)) { + final StringBuilder sb = new StringBuilder(); + buf.print(sb, /*indent*/ 0); + System.out.println(sb.toString()); + } + } + +.. code-block:: text + + ArrowBuf[2], address:140433199984656, length:1024 + event log for: ArrowBuf[2] + 675959093395667 create() + at org.apache.arrow.memory.util.HistoricalLog$Event.(HistoricalLog.java:175) + at org.apache.arrow.memory.util.HistoricalLog.recordEvent(HistoricalLog.java:83) + at org.apache.arrow.memory.ArrowBuf.(ArrowBuf.java:96) + at org.apache.arrow.memory.BufferLedger.newArrowBuf(BufferLedger.java:271) + at org.apache.arrow.memory.BaseAllocator.bufferWithoutReservation(BaseAllocator.java:300) + at org.apache.arrow.memory.BaseAllocator.buffer(BaseAllocator.java:276) + at org.apache.arrow.memory.RootAllocator.buffer(RootAllocator.java:29) + at org.apache.arrow.memory.BaseAllocator.buffer(BaseAllocator.java:240) + at org.apache.arrow.memory.RootAllocator.buffer(RootAllocator.java:29) + at REPL.$JShell$14.do_it$($JShell$14.java:10) + at jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(NativeMethodAccessorImpl.java:-2) + at jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62) + at jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43) + at java.lang.reflect.Method.invoke(Method.java:566) + at jdk.jshell.execution.DirectExecutionControl.invoke(DirectExecutionControl.java:209) + at jdk.jshell.execution.RemoteExecutionControl.invoke(RemoteExecutionControl.java:116) + at jdk.jshell.execution.DirectExecutionControl.invoke(DirectExecutionControl.java:119) + at jdk.jshell.execution.ExecutionControlForwarder.processCommand(ExecutionControlForwarder.java:144) + at jdk.jshell.execution.ExecutionControlForwarder.commandLoop(ExecutionControlForwarder.java:262) + at jdk.jshell.execution.Util.forwardExecutionControl(Util.java:76) + at jdk.jshell.execution.Util.forwardExecutionControlAndIO(Util.java:137) + at jdk.jshell.execution.RemoteExecutionControl.main(RemoteExecutionControl.java:70) + +The BufferAllocator also provides a ``BufferAllocator.toVerboseString()`` which can be used in +``DEBUG`` mode to get extensive stacktrace information and events associated with various Allocator behaviors. + +Finally, enabling the ``TRACE`` logging level will automatically provide this stack trace when the allocator is closed: + +.. code-block:: java + + // Assumes use of Logback; adjust for Log4j, etc. as appropriate + import ch.qos.logback.classic.Level; + import ch.qos.logback.classic.Logger; + import org.apache.arrow.memory.ArrowBuf; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.slf4j.LoggerFactory; + + // Set log level to TRACE to get tracebacks + ((Logger) LoggerFactory.getLogger("org.apache.arrow")).setLevel(Level.TRACE); + try (final BufferAllocator allocator = new RootAllocator()) { + // Leak buffer + allocator.buffer(1024); + } + +.. code-block:: text + + | Exception java.lang.IllegalStateException: Allocator[ROOT] closed with outstanding buffers allocated (1). + Allocator(ROOT) 0/1024/1024/9223372036854775807 (res/actual/peak/limit) + child allocators: 0 + ledgers: 1 + ledger[1] allocator: ROOT), isOwning: , size: , references: 1, life: 712040870231544..0, allocatorManager: [, life: ] holds 1 buffers. + ArrowBuf[2], address:139926571810832, length:1024 + event log for: ArrowBuf[2] + 712040888650134 create() + at org.apache.arrow.memory.util.StackTrace.(StackTrace.java:34) + at org.apache.arrow.memory.util.HistoricalLog$Event.(HistoricalLog.java:175) + at org.apache.arrow.memory.util.HistoricalLog.recordEvent(HistoricalLog.java:83) + at org.apache.arrow.memory.ArrowBuf.(ArrowBuf.java:96) + at org.apache.arrow.memory.BufferLedger.newArrowBuf(BufferLedger.java:271) + at org.apache.arrow.memory.BaseAllocator.bufferWithoutReservation(BaseAllocator.java:300) + at org.apache.arrow.memory.BaseAllocator.buffer(BaseAllocator.java:276) + at org.apache.arrow.memory.RootAllocator.buffer(RootAllocator.java:29) + at org.apache.arrow.memory.BaseAllocator.buffer(BaseAllocator.java:240) + at org.apache.arrow.memory.RootAllocator.buffer(RootAllocator.java:29) + at REPL.$JShell$18.do_it$($JShell$18.java:13) + at jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(NativeMethodAccessorImpl.java:-2) + at jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62) + at jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43) + at java.lang.reflect.Method.invoke(Method.java:566) + at jdk.jshell.execution.DirectExecutionControl.invoke(DirectExecutionControl.java:209) + at jdk.jshell.execution.RemoteExecutionControl.invoke(RemoteExecutionControl.java:116) + at jdk.jshell.execution.DirectExecutionControl.invoke(DirectExecutionControl.java:119) + at jdk.jshell.execution.ExecutionControlForwarder.processCommand(ExecutionControlForwarder.java:144) + at jdk.jshell.execution.ExecutionControlForwarder.commandLoop(ExecutionControlForwarder.java:262) + at jdk.jshell.execution.Util.forwardExecutionControl(Util.java:76) + at jdk.jshell.execution.Util.forwardExecutionControlAndIO(Util.java:137) + + reservations: 0 + + | at BaseAllocator.close (BaseAllocator.java:405) + | at RootAllocator.close (RootAllocator.java:29) + | at (#8:1) + +Sometimes, explicitly passing allocators around is difficult. For example, it +can be hard to pass around extra state, like an allocator, through layers of +existing application or framework code. A global or singleton allocator instance +can be useful here, though it should not be your first choice. + +How this works: + +1. Set up a global allocator in a singleton class. +2. Provide methods to create child allocators from the global allocator. +3. Give child allocators proper names to make it easier to figure out where + allocations occurred in case of errors. +4. Ensure that resources are properly closed. +5. Check that the global allocator is empty at some suitable point, such as + right before program shutdown. +6. If it is not empty, review the above allocation bugs. + +.. code-block:: java + + //1 + private static final BufferAllocator allocator = new RootAllocator(); + private static final AtomicInteger childNumber = new AtomicInteger(0); + ... + //2 + public static BufferAllocator getChildAllocator() { + return allocator.newChildAllocator(nextChildName(), 0, Long.MAX_VALUE); + } + ... + //3 + private static String nextChildName() { + return "Allocator-Child-" + childNumber.incrementAndGet(); + } + ... + //4: Business code + try (BufferAllocator allocator = GlobalAllocator.getChildAllocator()) { + ... + } + ... + //5 + public static void checkGlobalCleanUpResources() { + ... + if (!allocator.getChildAllocators().isEmpty()) { + throw new IllegalStateException(...); + } else if (allocator.getAllocatedMemory() != 0) { + throw new IllegalStateException(...); + } + } + +.. _`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/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 +===================== + +Design Principles +----------------- +Arrow’s memory model is based on the following basic concepts: + +- Memory can be allocated up to some limit. That limit could be a real + limit (OS/JVM) or a locally imposed limit. +- Allocation operates in two phases: accounting then actual allocation. + Allocation could fail at either point. +- Allocation failure should be recoverable. In all cases, the Allocator + infrastructure should expose memory allocation failures (OS or + internal limit-based) as ``OutOfMemoryException``\ s. +- Any allocator can reserve memory when created. This memory shall be + held such that this allocator will always be able to allocate that + amount of memory. +- A particular application component should work to use a local + allocator to understand local memory usage and better debug memory + leaks. +- The same physical memory can be shared by multiple allocators and the + allocator must provide an accounting paradigm for this purpose. + +Reserving Memory +---------------- + +Arrow provides two different ways to reserve memory: + +- BufferAllocator accounting reservations: When a new allocator (other + than the ``RootAllocator``) is initialized, it can set aside memory + that it will keep locally for its lifetime. This is memory that will + never be released back to its parent allocator until the allocator is + closed. +- ``AllocationReservation`` via BufferAllocator.newReservation(): + Allows a short-term preallocation strategy so that a particular + subsystem can ensure future memory is available to support a + particular request. + +Reference Counting Details +-------------------------- + +Typically, the ReferenceManager implementation used is an instance of `BufferLedger`_. +A BufferLedger is a ReferenceManager that also maintains the relationship between an ``AllocationManager``, +a ``BufferAllocator`` and one or more individual ``ArrowBuf``\ s + +All ArrowBufs (direct or sliced) related to a single BufferLedger/BufferAllocator combination +share the same reference count and either all will be valid or all will be invalid. +For simplicity of accounting, we treat that memory as being used by one +of the BufferAllocators associated with the memory. When that allocator +releases its claim on that memory, the memory ownership is then moved to +another BufferLedger belonging to the same AllocationManager. + +Allocation Details +------------------ + +There are several Allocator types in Arrow Java: + +- ``BufferAllocator`` - The public interface application users should be leveraging +- ``BaseAllocator`` - The base implementation of memory allocation, contains the meat of the Arrow allocator implementation +- ``RootAllocator`` - The root allocator. Typically only one created for a JVM. It serves as the parent/ancestor for child allocators +- ``ChildAllocator`` - A child allocator that derives from the root allocator + +Many BufferAllocators can reference the same piece of physical memory at the same +time. It is the AllocationManager’s responsibility to ensure that in this situation, +all memory is accurately accounted for from the Root’s perspective +and also to ensure that the memory is correctly released once all +BufferAllocators have stopped using that memory. + +For simplicity of accounting, we treat that memory as being used by one +of the BufferAllocators associated with the memory. When that allocator +releases its claim on that memory, the memory ownership is then moved to +another BufferLedger belonging to the same AllocationManager. Note that +because a ArrowBuf.release() is what actually causes memory ownership +transfer to occur, we always proceed with ownership transfer (even if +that violates an allocator limit). It is the responsibility of the +application owning a particular allocator to frequently confirm whether +the allocator is over its memory limit (BufferAllocator.isOverLimit()) +and if so, attempt to aggressively release memory to ameliorate the +situation. + + +Object Hierarchy +---------------- + +There are two main ways that someone can look at the object hierarchy +for Arrow’s memory management scheme. The first is a memory based +perspective as below: + +Memory Perspective +~~~~~~~~~~~~~~~~~~ + +.. code-block:: none + + + AllocationManager + | + |-- UnsignedDirectLittleEndian (One per AllocationManager) + | + |-+ BufferLedger 1 ==> Allocator A (owning) + | ` - ArrowBuf 1 + |-+ BufferLedger 2 ==> Allocator B (non-owning) + | ` - ArrowBuf 2 + |-+ BufferLedger 3 ==> Allocator C (non-owning) + | - ArrowBuf 3 + | - ArrowBuf 4 + ` - ArrowBuf 5 + +In this picture, a piece of memory is owned by an allocator manager. An +allocator manager is responsible for that piece of memory no matter +which allocator(s) it is working with. An allocator manager will have +relationships with a piece of raw memory (via its reference to +UnsignedDirectLittleEndian) as well as references to each +BufferAllocator it has a relationship to. + +Allocator Perspective +~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: none + + + RootAllocator + |-+ ChildAllocator 1 + | | - ChildAllocator 1.1 + | ` ... + | + |-+ ChildAllocator 2 + |-+ ChildAllocator 3 + | | + | |-+ BufferLedger 1 ==> AllocationManager 1 (owning) ==> UDLE + | | `- ArrowBuf 1 + | `-+ BufferLedger 2 ==> AllocationManager 2 (non-owning)==> UDLE + | `- ArrowBuf 2 + | + |-+ BufferLedger 3 ==> AllocationManager 1 (non-owning)==> UDLE + | ` - ArrowBuf 3 + |-+ BufferLedger 4 ==> AllocationManager 2 (owning) ==> UDLE + | - ArrowBuf 4 + | - ArrowBuf 5 + ` - ArrowBuf 6 + +In this picture, a RootAllocator owns three ChildAllocators. The first +ChildAllocator (ChildAllocator 1) owns a subsequent ChildAllocator. +ChildAllocator has two BufferLedgers/AllocationManager references. +Coincidentally, each of these AllocationManager’s is also associated +with the RootAllocator. In this case, one of the these +AllocationManagers is owned by ChildAllocator 3 (AllocationManager 1) +while the other AllocationManager (AllocationManager 2) is +owned/accounted for by the RootAllocator. Note that in this scenario, +ArrowBuf 1 is sharing the underlying memory as ArrowBuf 3. However the +subset of that memory (e.g. through slicing) might be different. Also +note that ArrowBuf 2 and ArrowBuf 4, 5 and 6 are also sharing the same +underlying memory. Also note that ArrowBuf 4, 5 and 6 all share the same +reference count and fate. diff --git a/docs/source/overview.rst b/docs/source/overview.rst new file mode 100644 index 0000000000..1188054114 --- /dev/null +++ b/docs/source/overview.rst @@ -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. + +=================== +High-Level Overview +=================== + +The Apache Arrow Java modules implement various specifications including the +columnar format and IPC. Most modules are native Java implementations, +but some modules are JNI bindings to the C++ library. + +.. list-table:: Arrow Java Modules + :widths: 25 50 25 + :header-rows: 1 + + * - Module + - Description + - Implementation + * - arrow-format + - Generated Java files from the IPC Flatbuffer definitions. + - Native + * - arrow-memory-core + - Core off-heap memory management libraries for Arrow ValueVectors. + - Native + * - arrow-memory-unsafe + - Memory management implementation based on sun.misc.Unsafe. + - Native + * - arrow-memory-netty + - Memory management implementation based on Netty. + - Native + * - 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 + * - arrow-jdbc + - (Experimental) A library for converting JDBC data to Arrow data. + - Native + * - flight-core + - An RPC mechanism for transferring ValueVectors. + - Native + * - flight-sql + - Contains utility classes to expose Flight SQL semantics for clients and servers over Arrow Flight. + - Native + * - flight-integration-tests + - Integration tests for Flight RPC. + - Native + * - arrow-performance + - JMH benchmarks for the Arrow libraries. + - Native + * - arrow-algorithm + - (Experimental) A collection of algorithms for working with ValueVectors. + - Native + * - arrow-avro + - (Experimental) A library for converting Avro data to Arrow data. + - Native + * - arrow-compression + - (Experimental) A library for working with compression/decompression of Arrow data. + - Native + * - arrow-c-data + - Java implementation of `C Data Interface`_ + - JNI + * - arrow-orc + - (Experimental) A JNI wrapper for the C++ ORC reader implementation. + - JNI + * - arrow-gandiva + - Java wrappers around the native Gandiva SQL expression compiler. + - JNI + * - arrow-dataset + - Java bindings to the Arrow Datasets library. + - JNI + +Arrow Java modules support working with data (1) in-memory, (2) at rest, and (3) on-the-wire. + +.. _`C Data Interface`: https://arrow.apache.org/docs/format/CDataInterface.html diff --git a/docs/source/quickstartguide.rst b/docs/source/quickstartguide.rst new file mode 100644 index 0000000000..adb07d7002 --- /dev/null +++ b/docs/source/quickstartguide.rst @@ -0,0 +1,314 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +================= +Quick Start Guide +================= + +Arrow Java provides several building blocks. Data types describe the types of values; +ValueVectors are sequences of typed values; fields describe the types of columns in +tabular data; schemas describe a sequence of columns in tabular data, and +VectorSchemaRoot represents tabular data. Arrow also provides readers and +writers for loading data from and persisting data to storage. + +Create a ValueVector +******************** + +**ValueVectors** represent a sequence of values of the same type. +They are also known as "arrays" in the columnar format. + +Example: create a vector of 32-bit integers representing ``[1, null, 2]``: + +.. code-block:: Java + + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.IntVector; + + try( + BufferAllocator allocator = new RootAllocator(); + IntVector intVector = new IntVector("fixed-size-primitive-layout", allocator); + ){ + intVector.allocateNew(3); + intVector.set(0,1); + intVector.setNull(1); + intVector.set(2,2); + intVector.setValueCount(3); + System.out.println("Vector created in memory: " + intVector); + } + +.. code-block:: shell + + Vector created in memory: [1, null, 2] + + +Example: create a vector of UTF-8 encoded strings representing ``["one", "two", "three"]``: + +.. code-block:: Java + + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.VarCharVector; + + try( + BufferAllocator allocator = new RootAllocator(); + VarCharVector varCharVector = new VarCharVector("variable-size-primitive-layout", allocator); + ){ + varCharVector.allocateNew(3); + varCharVector.set(0, "one".getBytes()); + varCharVector.set(1, "two".getBytes()); + varCharVector.set(2, "three".getBytes()); + varCharVector.setValueCount(3); + System.out.println("Vector created in memory: " + varCharVector); + } + +.. code-block:: shell + + Vector created in memory: [one, two, three] + +Create a Field +************** + +**Fields** are used to denote the particular columns of tabular data. +They consist of a name, a data type, a flag indicating whether the column can have null values, +and optional key-value metadata. + +Example: create a field named "document" of string type: + +.. code-block:: Java + + 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 java.util.HashMap; + import java.util.Map; + + Map metadata = new HashMap<>(); + metadata.put("A", "Id card"); + metadata.put("B", "Passport"); + metadata.put("C", "Visa"); + Field document = new Field("document", + new FieldType(true, new ArrowType.Utf8(), /*dictionary*/ null, metadata), + /*children*/ null); + System.out.println("Field created: " + document + ", Metadata: " + document.getMetadata()); + +.. code-block:: shell + + Field created: document: Utf8, Metadata: {A=Id card, B=Passport, C=Visa} + +Create a Schema +*************** + +**Schemas** hold a sequence of fields together with some optional metadata. + +Example: Create a schema describing datasets with two columns: +an int32 column "A" and a UTF8-encoded string column "B" + +.. code-block:: Java + + 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 java.util.HashMap; + import java.util.Map; + import static java.util.Arrays.asList; + + Map metadata = new HashMap<>(); + metadata.put("K1", "V1"); + metadata.put("K2", "V2"); + Field a = new Field("A", FieldType.nullable(new ArrowType.Int(32, true)), /*children*/ null); + Field b = new Field("B", FieldType.nullable(new ArrowType.Utf8()), /*children*/ null); + Schema schema = new Schema(asList(a, b), metadata); + System.out.println("Schema created: " + schema); + +.. code-block:: shell + + Schema created: Schema(metadata: {K1=V1, K2=V2}) + +Create a VectorSchemaRoot +************************* + +A **VectorSchemaRoot** combines ValueVectors with a Schema to represent tabular data. + +Example: Create a dataset of names (strings) and ages (32-bit signed integers). + +.. code-block:: Java + + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.IntVector; + import org.apache.arrow.vector.VarCharVector; + import org.apache.arrow.vector.VectorSchemaRoot; + 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 java.nio.charset.StandardCharsets; + import java.util.HashMap; + import java.util.Map; + import static java.util.Arrays.asList; + + Field age = new Field("age", + FieldType.nullable(new ArrowType.Int(32, true)), + /*children*/null + ); + Field name = new Field("name", + FieldType.nullable(new ArrowType.Utf8()), + /*children*/null + ); + Schema schema = new Schema(asList(age, name), /*metadata*/ null); + try( + BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + IntVector ageVector = (IntVector) root.getVector("age"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + ){ + ageVector.allocateNew(3); + ageVector.set(0, 10); + ageVector.set(1, 20); + ageVector.set(2, 30); + nameVector.allocateNew(3); + nameVector.set(0, "Dave".getBytes(StandardCharsets.UTF_8)); + nameVector.set(1, "Peter".getBytes(StandardCharsets.UTF_8)); + nameVector.set(2, "Mary".getBytes(StandardCharsets.UTF_8)); + root.setRowCount(3); + System.out.println("VectorSchemaRoot created: \n" + root.contentToTSVString()); + } + +.. code-block:: shell + + VectorSchemaRoot created: + age name + 10 Dave + 20 Peter + 30 Mary + + +Interprocess Communication (IPC) +******************************** + +Arrow data can be written to and read from disk, and both of these can be done in +a streaming and/or random-access fashion depending on application requirements. + +**Write data to an arrow file** + +Example: Write the dataset from the previous example to an Arrow IPC file (random-access). + +.. code-block:: Java + + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.IntVector; + import org.apache.arrow.vector.VarCharVector; + import org.apache.arrow.vector.VectorSchemaRoot; + import org.apache.arrow.vector.ipc.ArrowFileWriter; + 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 java.io.File; + import java.io.FileOutputStream; + import java.io.IOException; + import java.nio.charset.StandardCharsets; + import java.util.HashMap; + import java.util.Map; + import static java.util.Arrays.asList; + + Field age = new Field("age", + FieldType.nullable(new ArrowType.Int(32, true)), + /*children*/ null); + Field name = new Field("name", + FieldType.nullable(new ArrowType.Utf8()), + /*children*/ null); + Schema schema = new Schema(asList(age, name)); + try( + BufferAllocator allocator = new RootAllocator(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + IntVector ageVector = (IntVector) root.getVector("age"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + ){ + ageVector.allocateNew(3); + ageVector.set(0, 10); + ageVector.set(1, 20); + ageVector.set(2, 30); + nameVector.allocateNew(3); + nameVector.set(0, "Dave".getBytes(StandardCharsets.UTF_8)); + nameVector.set(1, "Peter".getBytes(StandardCharsets.UTF_8)); + nameVector.set(2, "Mary".getBytes(StandardCharsets.UTF_8)); + root.setRowCount(3); + File file = new File("random_access_file.arrow"); + try ( + FileOutputStream fileOutputStream = new FileOutputStream(file); + ArrowFileWriter writer = new ArrowFileWriter(root, /*provider*/ null, fileOutputStream.getChannel()); + ) { + writer.start(); + writer.writeBatch(); + writer.end(); + System.out.println("Record batches written: " + writer.getRecordBlocks().size() + + ". Number of rows written: " + root.getRowCount()); + } catch (IOException e) { + e.printStackTrace(); + } + } + +.. code-block:: shell + + Record batches written: 1. Number of rows written: 3 + +**Read data from an arrow file** + +Example: Read the dataset from the previous example from an Arrow IPC file (random-access). + +.. code-block:: Java + + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.ipc.ArrowFileReader; + import org.apache.arrow.vector.ipc.message.ArrowBlock; + import org.apache.arrow.vector.VectorSchemaRoot; + import java.io.File; + import java.io.FileInputStream; + import java.io.FileOutputStream; + import java.io.IOException; + + try( + BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + FileInputStream fileInputStream = new FileInputStream(new File("random_access_file.arrow")); + ArrowFileReader reader = new ArrowFileReader(fileInputStream.getChannel(), allocator); + ){ + System.out.println("Record batches in file: " + reader.getRecordBlocks().size()); + for (ArrowBlock arrowBlock : reader.getRecordBlocks()) { + reader.loadRecordBatch(arrowBlock); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + System.out.println("VectorSchemaRoot read: \n" + root.contentToTSVString()); + } + } catch (IOException e) { + e.printStackTrace(); + } + +.. code-block:: shell + + Record batches in file: 1 + VectorSchemaRoot read: + age name + 10 Dave + 20 Peter + 30 Mary + +More examples available at `Arrow Java Cookbook`_. + +.. _`Arrow Java Cookbook`: https://arrow.apache.org/cookbook/java diff --git a/docs/source/reference/index.rst b/docs/source/reference/index.rst new file mode 100644 index 0000000000..523ac0c7f7 --- /dev/null +++ b/docs/source/reference/index.rst @@ -0,0 +1,21 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +Java Reference (javadoc) +======================== + +Stub page for the Java reference docs; actual source is located in the java/ directory. diff --git a/docs/source/substrait.rst b/docs/source/substrait.rst new file mode 100644 index 0000000000..5ec07f1658 --- /dev/null +++ b/docs/source/substrait.rst @@ -0,0 +1,201 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +========= +Substrait +========= + +The ``arrow-dataset`` module can execute Substrait_ plans via the :external+arrow:doc:`Acero ` +query engine. + +Executing Queries Using Substrait Plans +======================================= + +Plans can reference data in files via URIs, or "named tables" that must be provided along with the plan. + +Here is an example of a Java program that queries a Parquet file using Java Substrait +(this example use `Substrait Java`_ project to compile a SQL query to a Substrait plan): + +.. code-block:: Java + + import com.google.common.collect.ImmutableList; + import io.substrait.isthmus.SqlToSubstrait; + import io.substrait.proto.Plan; + import org.apache.arrow.dataset.file.FileFormat; + import org.apache.arrow.dataset.file.FileSystemDatasetFactory; + import org.apache.arrow.dataset.jni.NativeMemoryPool; + import org.apache.arrow.dataset.scanner.ScanOptions; + import org.apache.arrow.dataset.scanner.Scanner; + import org.apache.arrow.dataset.source.Dataset; + import org.apache.arrow.dataset.source.DatasetFactory; + import org.apache.arrow.dataset.substrait.AceroSubstraitConsumer; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.ipc.ArrowReader; + import org.apache.calcite.sql.parser.SqlParseException; + + import java.nio.ByteBuffer; + import java.util.HashMap; + import java.util.Map; + + public class ClientSubstrait { + public static void main(String[] args) { + String uri = "file:///data/tpch_parquet/nation.parquet"; + ScanOptions options = new ScanOptions(/*batchSize*/ 32768); + try ( + BufferAllocator allocator = new RootAllocator(); + DatasetFactory datasetFactory = new FileSystemDatasetFactory(allocator, NativeMemoryPool.getDefault(), + FileFormat.PARQUET, uri); + Dataset dataset = datasetFactory.finish(); + Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches() + ) { + // map table to reader + Map mapTableToArrowReader = new HashMap<>(); + mapTableToArrowReader.put("NATION", reader); + // get binary plan + Plan plan = getPlan(); + ByteBuffer substraitPlan = ByteBuffer.allocateDirect(plan.toByteArray().length); + substraitPlan.put(plan.toByteArray()); + // run query + try (ArrowReader arrowReader = new AceroSubstraitConsumer(allocator).runQuery( + substraitPlan, + mapTableToArrowReader + )) { + while (arrowReader.loadNextBatch()) { + System.out.println(arrowReader.getVectorSchemaRoot().contentToTSVString()); + } + } + } catch (Exception e) { + e.printStackTrace(); + } + } + + static Plan getPlan() throws SqlParseException { + String sql = "SELECT * from nation"; + String nation = "CREATE TABLE NATION (N_NATIONKEY BIGINT NOT NULL, N_NAME CHAR(25), " + + "N_REGIONKEY BIGINT NOT NULL, N_COMMENT VARCHAR(152))"; + SqlToSubstrait sqlToSubstrait = new SqlToSubstrait(); + Plan plan = sqlToSubstrait.execute(sql, ImmutableList.of(nation)); + return plan; + } + } + +.. code-block:: text + + // Results example: + FieldPath(0) FieldPath(1) FieldPath(2) FieldPath(3) + 0 ALGERIA 0 haggle. carefully final deposits detect slyly agai + 1 ARGENTINA 1 al foxes promise slyly according to the regular accounts. bold requests alon + +Executing Projections and Filters Using Extended Expressions +============================================================ + +Dataset also supports projections and filters with Substrait's `Extended Expression`_. +This requires the substrait-java library. + +This Java program: + +- Loads a Parquet file containing the "nation" table from the TPC-H benchmark. +- Applies a filter: + - ``N_NATIONKEY > 18`` +- Projects two new columns: + - ``N_REGIONKEY + 10`` + - ``N_NAME || ' - ' || N_COMMENT`` + + + +.. code-block:: Java + + import com.google.common.collect.ImmutableList; + import io.substrait.isthmus.SqlExpressionToSubstrait; + import io.substrait.proto.ExtendedExpression; + import org.apache.arrow.dataset.file.FileFormat; + import org.apache.arrow.dataset.file.FileSystemDatasetFactory; + import org.apache.arrow.dataset.jni.NativeMemoryPool; + import org.apache.arrow.dataset.scanner.ScanOptions; + import org.apache.arrow.dataset.scanner.Scanner; + import org.apache.arrow.dataset.source.Dataset; + import org.apache.arrow.dataset.source.DatasetFactory; + import org.apache.arrow.memory.BufferAllocator; + import org.apache.arrow.memory.RootAllocator; + import org.apache.arrow.vector.ipc.ArrowReader; + import org.apache.calcite.sql.parser.SqlParseException; + + import java.nio.ByteBuffer; + import java.util.Base64; + import java.util.Optional; + + public class ClientSubstraitExtendedExpressionsCookbook { + + public static void main(String[] args) throws SqlParseException { + projectAndFilterDataset(); + } + + private static void projectAndFilterDataset() throws SqlParseException { + String uri = "file:///Users/data/tpch_parquet/nation.parquet"; + ScanOptions options = + new ScanOptions.Builder(/*batchSize*/ 32768) + .columns(Optional.empty()) + .substraitFilter(getByteBuffer(new String[]{"N_NATIONKEY > 18"})) + .substraitProjection(getByteBuffer(new String[]{"N_REGIONKEY + 10", + "N_NAME || CAST(' - ' as VARCHAR) || N_COMMENT"})) + .build(); + try (BufferAllocator allocator = new RootAllocator(); + DatasetFactory datasetFactory = + new FileSystemDatasetFactory( + allocator, NativeMemoryPool.getDefault(), FileFormat.PARQUET, uri); + Dataset dataset = datasetFactory.finish(); + Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + while (reader.loadNextBatch()) { + System.out.println(reader.getVectorSchemaRoot().contentToTSVString()); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static ByteBuffer getByteBuffer(String[] sqlExpression) throws SqlParseException { + String schema = + "CREATE TABLE NATION (N_NATIONKEY INT NOT NULL, N_NAME VARCHAR, " + + "N_REGIONKEY INT NOT NULL, N_COMMENT VARCHAR)"; + SqlExpressionToSubstrait expressionToSubstrait = new SqlExpressionToSubstrait(); + ExtendedExpression expression = + expressionToSubstrait.convert(sqlExpression, ImmutableList.of(schema)); + byte[] expressionToByte = + Base64.getDecoder().decode(Base64.getEncoder().encodeToString(expression.toByteArray())); + ByteBuffer byteBuffer = ByteBuffer.allocateDirect(expressionToByte.length); + byteBuffer.put(expressionToByte); + return byteBuffer; + } + } + +.. code-block:: text + + column-1 column-2 + 13 ROMANIA - ular asymptotes are about the furious multipliers. express dependencies nag above the ironically ironic account + 14 SAUDI ARABIA - ts. silent requests haggle. closely express packages sleep across the blithely + 12 VIETNAM - hely enticingly express accounts. even, final + 13 RUSSIA - requests against the platelets use never according to the quickly regular pint + 13 UNITED KINGDOM - eans boost carefully special requests. accounts are. carefull + 11 UNITED STATES - y final packages. slow foxes cajole quickly. quickly silent platelets breach ironic accounts. unusual pinto be + +.. _`Substrait`: https://substrait.io/ +.. _`Substrait Java`: https://github.com/substrait-io/substrait-java +.. _`Acero`: https://arrow.apache.org/docs/cpp/streaming_execution.html +.. _`Extended Expression`: https://github.com/substrait-io/substrait/blob/main/site/docs/expressions/extended_expression.md diff --git a/docs/source/table.rst b/docs/source/table.rst new file mode 100644 index 0000000000..880ef84d29 --- /dev/null +++ b/docs/source/table.rst @@ -0,0 +1,378 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +===== +Table +===== + +**NOTE**: The Table API is experimental and subject to change. See the list of limitations below. + +`Table`_ is an immutable tabular data structure based on `FieldVector`_. Like `VectorSchemaRoot`_, ``Table`` is a columnar data structure backed by Arrow arrays, or more specifically, by ``FieldVector`` objects. It differs from ``VectorSchemaRoot`` mainly in that it is fully immutable and lacks support for batch operations. Anyone processing batches of tabular data in a pipeline should continue to use ``VectorSchemaRoot``. Finally, the ``Table`` API is mainly row-oriented, so in some ways it's more like the JDBC API than the ``VectorSchemaRoot`` API, but you can still use ``FieldReaders`` to work with data in a columnar fashion. + +Mutation in Table and VectorSchemaRoot +====================================== + +``VectorSchemaRoot`` provides a thin wrapper on the vectors that hold its data. Individual vectors can be retrieved from a vector schema root. These vectors have *setters* for modifying their elements, making ``VectorSchemaRoot`` immutable only by convention. The protocol for mutating a vector is documented in the `ValueVector`_ interface: + +- values need to be written in order (e.g. index 0, 1, 2, 5) +- null vectors start with all values as null before writing anything +- for variable width types, the offset vector should be all zeros before writing +- you must call setValueCount before a vector can be read +- you should never write to a vector once it has been read. + +The rules aren't enforced by the API so the programmer is responsible for ensuring that they are followed. Failure to do so could lead to runtime exceptions. + +``Table``, on the other hand, is immutable. The underlying vectors are not exposed. When a table is created from existing vectors, their memory is transferred to new vectors, so subsequent changes to the original vectors can't impact the new table's values. + +Features and limitations +====================================== + +A basic set of table functionality is currently available: + +- Create a table from vectors or ``VectorSchemaRoot`` +- Iterate tables by row, or set the current row index directly +- Access vector values as primitives, objects, and/or nullable `ValueHolder`_ instances (depending on type) +- Get a ``FieldReader`` for any vector +- Add and remove vectors, creating new tables +- Encode and decode a table's vectors using dictionary encoding +- Export table data for use by native code +- Print representative data to TSV strings +- Get a table's schema +- Slice tables +- Convert table to ``VectorSchemaRoot`` + +Limitations in the 11.0.0 release: + +- No support ``ChunkedArray`` or any form of row-group. Support for chunked arrays or row groups will be considered for a future release. +- No support for the C-Stream API. Support for the streaming API is contingent on chunked array support +- No support for creating tables directly from Java POJOs. All data held by a table must be imported via a ``VectorSchemaRoot``, or from collections or arrays of vectors. + +The Table API +============= + +Like ``VectorSchemaRoot``, a table contains a `Schema`_ and an ordered collection of ``FieldVector`` objects, but it is designed to be accessed via a row-oriented interface. + +Creating a Table from a VectorSchemaRoot +**************************************** + +Tables are created from a ``VectorSchemaRoot`` as shown below. The memory buffers holding the data are transferred from the vector schema root to new vectors in the new table, clearing the source vectors in the process. This ensures that the data in your new table is never changed. Since the buffers are transferred rather than copied, this is a very low overhead operation. + +.. code-block:: Java + + Table t = new Table(someVectorSchemaRoot); + +If you now update the vectors held by the ``VectorSchemaRoot`` (using some version of ``ValueVector#setSafe()``), it would reflect those changes, but the values in table *t* are unchanged. + +Creating a Table from FieldVectors +********************************** + +Tables can be created from ``FieldVectors`` as shown below, using 'var-arg' array arguments: + +.. code-block:: Java + + IntVector myVector = createMyIntVector(); + VectorSchemaRoot vsr1 = new VectorSchemaRoot(myVector); + +or by passing a collection: + +.. code-block:: Java + + IntVector myVector = createMyIntVector(); + List fvList = List.of(myVector); + VectorSchemaRoot vsr1 = new VectorSchemaRoot(fvList); + +It is rarely a good idea to share vectors between multiple vector schema roots, and it would not be a good idea to share them between vector schema roots and tables. Creating a ``VectorSchemaRoot`` from a list of vectors does not cause the reference counts for the vectors to be incremented. Unless you manage the counts manually, the code below would lead to more references than reference counts, and that could lead to trouble. There is an implicit assumption that the vectors were created for use by *one* ``VectorSchemaRoot`` that this code violates. + +*Don't do this:* + +.. code-block:: Java + + IntVector myVector = createMyIntVector(); // Reference count for myVector = 1 + VectorSchemaRoot vsr1 = new VectorSchemaRoot(myVector); // Still one reference + VectorSchemaRoot vsr2 = new VectorSchemaRoot(myVector); + // Ref count is still one, but there are two VSRs with a reference to myVector + vsr2.clear(); // Reference count for myVector is 0. + +What is happening is that the reference counter works at a lower level than the ``VectorSchemaRoot`` interface. A reference counter counts references to `ArrowBuf`_ instances that control memory buffers. It doesn't count references to the vectors that hold those ArrowBufs. In the example above, each ``ArrowBuf`` is held by one vector, so there is only one reference. This distinction is blurred when you call the ``VectorSchemaRoot``'s clear() method, which frees the memory held by each of the vectors it references even though another instance references the same vectors. + +When you create tables from vectors, it's assumed that there are no external references to those vectors. To be certain, the buffers underlying these vectors are transferred to new vectors in the new table, and the original vectors are cleared. + +*Don't do this either, but note the difference from above:* + +.. code-block:: Java + + IntVector myVector = createMyIntVector(); // Reference count for myVector = 1 + Table t1 = new Table(myVector); + // myVector is cleared; Table t1 has a new hidden vector with the data from myVector + Table t2 = new Table(myVector); + // t2 has no rows because myVector was just cleared + // t1 continues to have the data from the original vector + t2.clear(); + // no change because t2 is already empty and t1 is independent + +With tables, memory is explicitly transferred on instantiation so the buffers held by a table are held by *only* that table. + +Creating Tables with dictionary-encoded vectors +*********************************************** + +Another point of difference is that ``VectorSchemaRoot`` is uninformed about any dictionary-encoding of its vectors, while tables hold an optional `DictionaryProvider`_ instance. If any vectors in the source data are encoded, a DictionaryProvider must be set to un-encode the values. + +.. code-block:: Java + + VectorSchemaRoot vsr = myVsr(); + DictionaryProvider provider = myProvider(); + Table t = new Table(vsr, provider); + +In ``Table``, dictionaries are used like they are with vectors. To decode a vector, the user provides the name of the vector to decode and the dictionary id: + +.. code-block:: Java + + Table t = new Table(vsr, provider); + ValueVector decodedName = t.decode("name", 1L); + +To encode a vector from a table, a similar approach is used: + +.. code-block:: Java + + Table t = new Table(vsr, provider); + ValueVector encodedName = t.encode("name", 1L); + +Freeing memory explicitly +************************* + +Tables use off-heap memory that must be freed when it is no longer needed. ``Table`` implements ``AutoCloseable`` so the best way to create one is in a try-with-resources block: + +.. code-block:: Java + + try (VectorSchemaRoot vsr = myMethodForGettingVsrs(); + Table t = new Table(vsr)) { + // do useful things. + } + +If you don't use a try-with-resources block, you must close the table manually: + +.. code-block:: Java + + try { + VectorSchemaRoot vsr = myMethodForGettingVsrs(); + Table t = new Table(vsr); + // do useful things. + } finally { + vsr.close(); + t.close(); + } + +Manual closing should be performed in a finally block. + +Getting the schema +****************** + +You get the table's schema just as you would with a vector schema root: + +.. code-block:: Java + + Schema s = table.getSchema(); + +Adding and removing vectors +*************************** + +``Table`` provides facilities for adding and removing vectors modeled on the same functionality in ``VectorSchemaRoot``. These operations return new instances rather than modifying the original instance in-place. + +.. code-block:: Java + + try (Table t = new Table(vectorList)) { + IntVector v3 = new IntVector("3", intFieldType, allocator); + Table t2 = t.addVector(2, v3); + Table t3 = t2.removeVector(1); + // don't forget to close t2 and t3 + } + +Slicing tables +************** + +``Table`` supports *slice()* operations, where a slice of a source table is a second Table that refers to a single, contiguous range of rows in the source. + +.. code-block:: Java + + try (Table t = new Table(vectorList)) { + Table t2 = t.slice(100, 200); // creates a slice referencing the values in range (100, 200] + ... + } + +This raises the question: If you create a slice with *all* the values in the source table (as shown below), how would that differ from a new Table constructed with the same vectors as the source? + +.. code-block:: Java + + try (Table t = new Table(vectorList)) { + Table t2 = t.slice(0, t.getRowCount()); // creates a slice referencing all the values in t + // ... + } + +The difference is that when you *construct* a new table, the buffers are transferred from the source vectors to new vectors in the destination. With a slice, both tables share the same underlying vectors. That's OK, though, since both tables are immutable. + +Using FieldReaders +****************** + +You can get a `FieldReader`_ for any vector in the Table passing either the `Field`_, vector index, or vector name as an argument. The signatures are the same as in ``VectorSchemaRoot``. + +.. code-block:: Java + + FieldReader nameReader = table.getReader("user_name"); + +Row operations +************** + +Row-based access is supported by the `Row`_ object. ``Row`` provides *get()* methods by both vector name and vector position, but no *set()* operations. + +It is important to recognize that rows are NOT reified as objects, but rather operate like a cursor where the data from numerous logical rows in the table can be viewed (one at a time) using the same ``Row`` instance. See "Moving from row-to-row" below for information about navigating through the table. + +Getting a row +************* + +Calling ``immutableRow()`` on any table instance returns a new ``Row`` instance. + +.. code-block:: Java + + Row r = table.immutableRow(); + +Moving from row-to-row +********************** + +Since rows are iterable, you can traverse a table using a standard while loop: + +.. code-block:: Java + + Row r = table.immutableRow(); + while (r.hasNext()) { + r.next(); + // do something useful here + } + +``Table`` implements ``Iterable`` so you can access rows directly from a table in an enhanced *for* loop: + +.. code-block:: Java + + for (Row row: table) { + int age = row.getInt("age"); + boolean nameIsNull = row.isNull("name"); + ... + } + +Finally, while rows are usually iterated in the order of the underlying data vectors, but they are also positionable using the ``Row#setPosition()`` method, so you can skip to a specific row. Row numbers are 0-based. + +.. code-block:: Java + + Row r = table.immutableRow(); + int age101 = r.setPosition(101); // change position directly to 101 + +Any changes to position are applied to all the columns in the table. + +Note that you must call ``next()``, or ``setPosition()`` before accessing values via a row. Failure to do so results in a runtime exception. + +Read operations using rows +************************** + +Methods are available for getting values by vector name and vector index, where index is the 0-based position of the vector in the table. For example, assuming 'age' is the 13th vector in 'table', the following two gets are equivalent: + +.. code-block:: Java + + Row r = table.immutableRow(); + r.next(); // position the row at the first value + int age1 = r.get("age"); // gets the value of vector named 'age' in the table at row 0 + int age2 = r.get(12); // gets the value of the 13th vector in the table at row 0 + +You can also get value using a nullable ``ValueHolder``. For example: + +.. code-block:: Java + + NullableIntHolder holder = new NullableIntHolder(); + int b = row.getInt("age", holder); + +This can be used to retrieve values without creating a new Object for each. + +In addition to getting values, you can check if a value is null using ``isNull()``. This is important if the vector contains any nulls, as asking for a value from a vector can cause NullPointerExceptions in some cases. + +.. code-block:: Java + + boolean name0isNull = row.isNull("name"); + +You can also get the current row number: + +.. code-block:: Java + + int row = row.getRowNumber(); + +Reading values as Objects +************************* + +For any given vector type, the basic *get()* method returns a primitive value wherever possible. For example, *getTimeStampMicro()* returns a long value that encodes the timestamp. To get the LocalDateTime object representing that timestamp in Java, another method with 'Obj' appended to the name is provided. For example: + +.. code-block:: Java + + long ts = row.getTimeStampMicro(); + LocalDateTime tsObject = row.getTimeStampMicroObj(); + +The exception to this naming scheme is for complex vector types (List, Map, Schema, Union, DenseUnion, and ExtensionType). These always return objects rather than primitives so no "Obj" extension is required. It is expected that some users may subclass ``Row`` to add getters that are more specific to their needs. + +Reading VarChars and LargeVarChars +********************************** + +Strings in arrow are represented as byte arrays encoded with the UTF-8 charset. You can get either a String result or the actual byte array. + +.. code-block:: Java + + byte[] b = row.getVarChar("first_name"); + String s = row.getVarCharObj("first_name"); // uses the default encoding (UTF-8) + +Converting a Table to a VectorSchemaRoot +**************************************** + +Tables can be converted to vector schema roots using the *toVectorSchemaRoot()* method. Buffers are transferred to the vector schema root and the source table is cleared. + +.. code-block:: Java + + VectorSchemaRoot root = myTable.toVectorSchemaRoot(); + +Working with the C-Data interface +********************************* + +The ability to work with native code is required for many Arrow features. This section describes how tables can be be exported for use with native code + +Exporting works by converting the data to a ``VectorSchemaRoot`` instance and using the existing facilities to transfer the data. You could do it yourself, but that isn't ideal because conversion to a vector schema root breaks the immutability guarantees. Using the ``exportTable()`` methods in the `Data`_ class avoids this concern. + +.. code-block:: Java + + Data.exportTable(bufferAllocator, table, dictionaryProvider, outArrowArray); + +If the table contains dictionary-encoded vectors and was constructed with a ``DictionaryProvider``, the provider argument to ``exportTable()`` can be omitted and the table's provider attribute will be used: + +.. code-block:: Java + + Data.exportTable(bufferAllocator, table, outArrowArray); + +.. _`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.rst b/docs/source/vector.rst new file mode 100644 index 0000000000..1996277444 --- /dev/null +++ b/docs/source/vector.rst @@ -0,0 +1,366 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +=========== +ValueVector +=========== + +:class:`ValueVector` interface (which called Array in C++ implementation and +the :external+arrow:doc:`the specification `) is an abstraction that is used to store a +sequence of values having the same type in an individual column. Internally, those values are +represented by one or several buffers, the number and meaning of which depend on the vector’s data type. + +There are concrete subclasses of :class:`ValueVector` for each primitive data type +and nested type described in the specification. There are a few differences in naming +with the type names described in the specification: +Table with non-intuitive names (BigInt = 64 bit integer, etc). + +It is important that vector is allocated before attempting to read or write, +:class:`ValueVector` "should" strive to guarantee this order of operation: +create > allocate > mutate > set value count > access > clear (or allocate to start the process over). +We will go through a concrete example to demonstrate each operation in the next section. + +Vector Life Cycle +================= + +As discussed above, each vector goes through several steps in its life cycle, +and each step is triggered by a vector operation. In particular, we have the following vector operations: + +1. **Vector creation**: we create a new vector object by, for example, the vector constructor. +The following code creates a new ``IntVector`` by the constructor: + +.. code-block:: Java + + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ... + IntVector vector = new IntVector("int vector", allocator); + +By now, a vector object is created. However, no underlying memory has been allocated, so we need the +following step. + +2. **Vector allocation**: in this step, we allocate memory for the vector. For most vectors, we +have two options: 1) if we know the maximum vector capacity, we can specify it by calling the +``allocateNew(int)`` method; 2) otherwise, we should call the ``allocateNew()`` method, and a default +capacity will be allocated for it. For our running example, we assume that the vector capacity never +exceeds 10: + +.. code-block:: Java + + vector.allocateNew(10); + +3. **Vector mutation**: now we can populate the vector with values we desire. For all vectors, we can populate +vector values through vector writers (An example will be given in the next section). For primitive types, +we can also mutate the vector by the set methods. There are two classes of set methods: 1) if we can +be sure the vector has enough capacity, we can call the ``set(index, value)`` method. 2) if we are not sure +about the vector capacity, we should call the ``setSafe(index, value)`` method, which will automatically +take care of vector reallocation, if the capacity is not sufficient. For our running example, we know the +vector has enough capacity, so we can call + +.. code-block:: Java + + vector.set(/*index*/5, /*value*/25); + +4. **Set value count**: for this step, we set the value count of the vector by calling the +``setValueCount(int)`` method: + +.. code-block:: Java + + vector.setValueCount(10); + +After this step, the vector enters an immutable state. In other words, we should no longer mutate it. +(Unless we reuse the vector by allocating it again. This will be discussed shortly.) + +5. **Vector access**: it is time to access vector values. Similarly, we have two options to access values: +1) get methods and 2) vector reader. Vector reader works for all types of vectors, while get methods are +only available for primitive vectors. A concrete example for vector reader will be given in the next section. +Below is an example of vector access by get method: + +.. code-block:: Java + + int value = vector.get(5); // value == 25 + +6. **Vector clear**: when we are done with the vector, we should clear it to release its memory. This is done by +calling the ``close()`` method: + +.. code-block:: Java + + vector.close(); + +Some points to note about the steps above: + +* The steps are not necessarily performed in a linear sequence. Instead, they can be in a loop. For example, + when a vector enters the access step, we can also go back to the vector mutation step, and then set value + count, access vector, and so on. + +* We should try to make sure the above steps are carried out in order. Otherwise, the vector + may be in an undefined state, and some unexpected behavior may occur. However, this restriction + is not strict. That means it is possible that we violates the order above, but still get + correct results. + +* When mutating vector values through set methods, we should prefer ``set(index, value)`` methods to + ``setSafe(index, value)`` methods whenever possible, to avoid unnecessary performance overhead of handling + vector capacity. + +* All vectors implement the ``AutoCloseable`` interface. So they must be closed explicitly when they are + no longer used, to avoid resource leak. To make sure of this, it is recommended to place vector related operations + into a try-with-resources block. + +* For fixed width vectors (e.g. IntVector), we can set values at different indices in arbitrary orders. + For variable width vectors (e.g. VarCharVector), however, we must set values in non-decreasing order of the + indices. Otherwise, the values after the set position will become invalid. For example, suppose we use the + following statements to populate a variable width vector: + +.. code-block:: Java + + VarCharVector vector = new VarCharVector("vector", allocator); + vector.allocateNew(); + vector.setSafe(0, "zero"); + vector.setSafe(1, "one"); + ... + vector.setSafe(9, "nine"); + +Then we set the value at position 5 again: + +.. code-block:: Java + + vector.setSafe(5, "5"); + +After that, the values at positions 6, 7, 8, and 9 of the vector will become invalid. + +Building ValueVector +==================== + +Note that the current implementation doesn't enforce the rule that Arrow objects are immutable. +:class:`ValueVector` instances could be created directly by using new keyword, there are +set/setSafe APIs and concrete subclasses of FieldWriter for populating values. + +For example, the code below shows how to build a :class:`BigIntVector`, in this case, we build a +vector of the range 0 to 7 where the element that should hold the fourth value is nulled + +.. code-block:: Java + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + BigIntVector vector = new BigIntVector("vector", allocator)) { + vector.allocateNew(8); + vector.set(0, 1); + vector.set(1, 2); + vector.set(2, 3); + vector.setNull(3); + vector.set(4, 5); + vector.set(5, 6); + vector.set(6, 7); + vector.set(7, 8); + vector.setValueCount(8); // this will finalizes the vector by convention. + ... + } + +The :class:`BigIntVector` holds two ArrowBufs. The first buffer holds the null bitmap, which consists +here of a single byte with the bits 1|1|1|1|0|1|1|1 (the bit is 1 if the value is non-null). +The second buffer contains all the above values. As the fourth entry is null, the value at that position +in the buffer is undefined. Note compared with set API, setSafe API would check value capacity before setting +values and reallocate buffers if necessary. + +Here is how to build a vector using writer + +.. code-block:: Java + + try (BigIntVector vector = new BigIntVector("vector", allocator); + BigIntWriter writer = new BigIntWriterImpl(vector)) { + writer.setPosition(0); + writer.writeBigInt(1); + writer.setPosition(1); + writer.writeBigInt(2); + writer.setPosition(2); + writer.writeBigInt(3); + // writer.setPosition(3) is not called which means the fourth value is null. + writer.setPosition(4); + writer.writeBigInt(5); + writer.setPosition(5); + writer.writeBigInt(6); + writer.setPosition(6); + writer.writeBigInt(7); + writer.setPosition(7); + writer.writeBigInt(8); + } + +There are get API and concrete subclasses of :class:`FieldReader` for accessing vector values, what needs +to be declared is that writer/reader is not as efficient as direct access + +.. code-block:: Java + + // access via get API + for (int i = 0; i < vector.getValueCount(); i++) { + if (!vector.isNull(i)) { + System.out.println(vector.get(i)); + } + } + + // access via reader + BigIntReader reader = vector.getReader(); + for (int i = 0; i < vector.getValueCount(); i++) { + reader.setPosition(i); + if (reader.isSet()) { + System.out.println(reader.readLong()); + } + } + +Building ListVector +=================== + +A :class:`ListVector` is a vector that holds a list of values for each index. Working with one you need to handle the same steps as mentioned above (create > allocate > mutate > set value count > access > clear), but the details of how you accomplish this are slightly different since you need to both create the vector and set the list of values for each index. + +For example, the code below shows how to build a :class:`ListVector` of int's using the writer :class:`UnionListWriter`. We build a vector from 0 to 9 and each index contains a list with values [[0, 0, 0, 0, 0], [0, 1, 2, 3, 4], [0, 2, 4, 6, 8], …, [0, 9, 18, 27, 36]]. List values can be added in any order so writing a list such as [3, 1, 2] would be just as valid. + +.. code-block:: Java + + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ListVector listVector = ListVector.empty("vector", allocator)) { + UnionListWriter writer = listVector.getWriter(); + for (int i = 0; i < 10; i++) { + writer.startList(); + writer.setPosition(i); + for (int j = 0; j < 5; j++) { + writer.writeInt(j * i); + } + writer.setValueCount(5); + writer.endList(); + } + listVector.setValueCount(10); + } + +:class:`ListVector` values can be accessed either through the get API or through the reader class :class:`UnionListReader`. To read all the values, first enumerate through the indexes, and then enumerate through the inner list values. + +.. code-block:: Java + + // access via get API + for (int i = 0; i < listVector.getValueCount(); i++) { + if (!listVector.isNull(i)) { + ArrayList elements = (ArrayList) listVector.getObject(i); + for (Integer element : elements) { + System.out.println(element); + } + } + } + + // access via reader + UnionListReader reader = listVector.getReader(); + for (int i = 0; i < listVector.getValueCount(); i++) { + reader.setPosition(i); + while (reader.next()) { + IntReader intReader = reader.reader(); + if (intReader.isSet()) { + System.out.println(intReader.readInteger()); + } + } + } + +Dictionary Encoding +=================== + +Dictionary encoding is a form of compression where values of one type are replaced by values of a smaller type: an array of ints replacing an array of strings is a common example. The mapping between the original values and the replacements is held in a 'dictionary'. Since the dictionary needs only one copy of each of the longer values, the combination of the dictionary and the array of smaller values may use less memory. The more repetitive the original data, the greater the savings. + +A ``FieldVector`` can be dictionary encoded for performance or improved memory efficiency. Nearly any type of vector might be encoded if there are many values, but few unique values. + +There are a few steps involved in the encoding process: + +1. Create a regular, un-encoded vector and populate it +2. Create a dictionary vector of the same type as the un-encoded vector. This vector must have the same values, but each unique value in the un-encoded vector need appear here only once. +3. Create a ``Dictionary``. It will contain the dictionary vector, plus a ``DictionaryEncoding`` object that holds the encoding's metadata and settings values. +4. Create a ``DictionaryEncoder``. +5. Call the encode() method on the ``DictionaryEncoder`` to produce an encoded version of the original vector. +6. (Optional) Call the decode() method on the encoded vector to re-create the original values. + +The encoded values will be integers. Depending on how many unique values you have, you can use ``TinyIntVector``, ``SmallIntVector``, ``IntVector``, or ``BigIntVector`` to hold them. You specify the type when you create your ``DictionaryEncoding`` instance. You might wonder where those integers come from: the dictionary vector is a regular vector, so the value's index position in that vector is used as its encoded value. + +Another critical attribute in ``DictionaryEncoding`` is the id. It's important to understand how the id is used, so we cover that later in this section. + +This result will be a new vector (for example, an ``IntVector``) that can act in place of the original vector (for example, a ``VarCharVector``). When you write the data in arrow format, it is both the new ``IntVector`` plus the dictionary that is written: you will need the dictionary later to retrieve the original values. + +.. code-block:: Java + + // 1. create a vector for the un-encoded data and populate it + VarCharVector unencoded = new VarCharVector("unencoded", allocator); + // now put some data in it before continuing + + // 2. create a vector to hold the dictionary and populate it + VarCharVector dictionaryVector = new VarCharVector("dictionary", allocator); + + // 3. create a dictionary object + Dictionary dictionary = new Dictionary(dictionaryVector, new DictionaryEncoding(1L, false, null)); + + // 4. create a dictionary encoder + DictionaryEncoder encoder = new DictionaryEncoder.encode(dictionary, allocator); + + // 5. encode the data + IntVector encoded = (IntVector) encoder.encode(unencoded); + + // 6. re-create an un-encoded version from the encoded vector + VarCharVector decoded = (VarCharVector) encoder.decode(encoded); + +One thing we haven't discussed is how to create the dictionary vector from the original un-encoded values. That is left to the library user since a custom method will likely be more efficient than a general utility. Since the dictionary vector is just a normal vector, you can populate its values with the standard APIs. + +Finally, you can package a number of dictionaries together, which is useful if you're working with a ``VectorSchemaRoot`` with several dictionary-encoded vectors. This is done using an object called a ``DictionaryProvider``. as shown in the example below. Note that we don't put the dictionary vectors in the same ``VectorSchemaRoot`` as the data vectors, as they will generally have fewer values. + + +.. code-block:: Java + + DictionaryProvider.MapDictionaryProvider provider = + new DictionaryProvider.MapDictionaryProvider(); + + provider.put(dictionary); + +The ``DictionaryProvider`` is simply a map of identifiers to ``Dictionary`` objects, where each identifier is a long value. In the above code you will see it as the first argument to the ``DictionaryEncoding`` constructor. + +This is where the ``DictionaryEncoding``'s 'id' attribute comes in. This value is used to connect dictionaries to instances of ``VectorSchemaRoot``, using a ``DictionaryProvider``. Here's how that works: + +* The ``VectorSchemaRoot`` has a ``Schema`` object containing a list of ``Field`` objects. +* The field has an attribute called 'dictionary', but it holds a ``DictionaryEncoding`` rather than a ``Dictionary`` +* As mentioned, the ``DictionaryProvider`` holds dictionaries indexed by a long value. This value is the id from your ``DictionaryEncoding``. +* To retrieve the dictionary for a vector in a ``VectorSchemaRoot``, you get the field associated with the vector, get its dictionary attribute, and use that object's id to look up the correct dictionary in the provider. + +.. code-block:: Java + + // create the encoded vector, the Dictionary and DictionaryProvider as discussed above + + // Create a VectorSchemaRoot with one encoded vector + VectorSchemaRoot vsr = new VectorSchemaRoot(List.of(encoded)); + + // now we want to decode our vector, so we retrieve its dictionary from the provider + Field f = vsr.getField(encoded.getName()); + DictionaryEncoding encoding = f.getDictionary(); + Dictionary dictionary = provider.lookup(encoding.getId()); + +As you can see, a ``DictionaryProvider`` is handy for managing the dictionaries associated with a ``VectorSchemaRoot``. More importantly, it helps package the dictionaries for a ``VectorSchemaRoot`` when it's written. The classes ``ArrowFileWriter`` and ``ArrowStreamWriter`` both accept an optional ``DictionaryProvider`` argument for that purpose. You can find example code for writing dictionaries in the documentation for (:doc:`ipc`). ``ArrowReader`` and its subclasses also implement the ``DictionaryProvider`` interface, so you can retrieve the actual dictionaries when reading a file. + +Slicing +======= + +Similar with C++ implementation, it is possible to make zero-copy slices of vectors to obtain a vector +referring to some logical sub-sequence of the data through :class:`TransferPair` + +.. code-block:: Java + + IntVector vector = new IntVector("intVector", allocator); + for (int i = 0; i < 10; i++) { + vector.setSafe(i, i); + } + vector.setValueCount(10); + + TransferPair tp = vector.getTransferPair(allocator); + tp.splitAndTransfer(0, 5); + IntVector sliced = (IntVector) tp.getTo(); + // In this case, the vector values are [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] and the sliceVector values are [0, 1, 2, 3, 4]. diff --git a/docs/source/vector_schema_root.rst b/docs/source/vector_schema_root.rst new file mode 100644 index 0000000000..f4a497c4e5 --- /dev/null +++ b/docs/source/vector_schema_root.rst @@ -0,0 +1,163 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +============ +Tabular Data +============ + +While arrays (aka: :doc:`ValueVector <./vector>`) represent a one-dimensional sequence of +homogeneous values, data often comes in the form of two-dimensional sets of +heterogeneous data (such as database tables, CSV files...). Arrow provides +several abstractions to handle such data conveniently and efficiently. + +Fields +====== + +Fields are used to denote the particular columns of tabular data. +A field, i.e. an instance of `Field`_, holds together a field name, a data +type, and some optional key-value metadata. + +.. code-block:: Java + + // Create a column "document" of string type with metadata + import org.apache.arrow.vector.types.pojo.ArrowType; + import org.apache.arrow.vector.types.pojo.Field; + import org.apache.arrow.vector.types.pojo.FieldType; + + Map metadata = new HashMap<>(); + metadata.put("A", "Id card"); + metadata.put("B", "Passport"); + metadata.put("C", "Visa"); + Field document = new Field("document", new FieldType(true, new ArrowType.Utf8(), /*dictionary*/ null, metadata), /*children*/ null); + +Schemas +======= + +A `Schema`_ describes the overall structure consisting of any number of columns. It holds a sequence of fields together +with some optional schema-wide metadata (in addition to per-field metadata). + +.. code-block:: Java + + // Create a schema describing datasets with two columns: + // a int32 column "A" and a utf8-encoded string column "B" + 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 static java.util.Arrays.asList; + + Map metadata = new HashMap<>(); + metadata.put("K1", "V1"); + metadata.put("K2", "V2"); + Field a = new Field("A", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field b = new Field("B", FieldType.nullable(new ArrowType.Utf8()), null); + Schema schema = new Schema(asList(a, b), metadata); + +VectorSchemaRoot +================ + +A `VectorSchemaRoot`_ is a container for batches of data. Batches flow through +VectorSchemaRoot as part of a pipeline. + +.. note:: + + VectorSchemaRoot is somewhat analogous to tables or record batches in the + other Arrow implementations in that they all are 2D datasets, but their + usage is different. + +The recommended usage is to create a single VectorSchemaRoot based on a known +schema and populate data over and over into that root in a stream of batches, +rather than creating a new instance each time (see `Flight`_ or +``ArrowFileWriter`` as examples). Thus at any one point, a VectorSchemaRoot may +have data or may have no data (say it was transferred downstream or not yet +populated). + +Here is an example of creating a VectorSchemaRoot: + +.. code-block:: Java + + BitVector bitVector = new BitVector("boolean", allocator); + VarCharVector varCharVector = new VarCharVector("varchar", allocator); + bitVector.allocateNew(); + varCharVector.allocateNew(); + for (int i = 0; i < 10; i++) { + bitVector.setSafe(i, i % 2 == 0 ? 0 : 1); + varCharVector.setSafe(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); + } + bitVector.setValueCount(10); + varCharVector.setValueCount(10); + + List fields = Arrays.asList(bitVector.getField(), varCharVector.getField()); + List vectors = Arrays.asList(bitVector, varCharVector); + VectorSchemaRoot vectorSchemaRoot = new VectorSchemaRoot(fields, vectors); + +Data can be loaded into/unloaded from a VectorSchemaRoot via `VectorLoader`_ +and `VectorUnloader`_. They handle converting between VectorSchemaRoot and +`ArrowRecordBatch`_ (a representation of a RecordBatch +:external+arrow:ref:`IPC ` message). For example: + +.. code-block:: Java + + // create a VectorSchemaRoot root1 and convert its data into recordBatch + VectorSchemaRoot root1 = new VectorSchemaRoot(fields, vectors); + VectorUnloader unloader = new VectorUnloader(root1); + ArrowRecordBatch recordBatch = unloader.getRecordBatch(); + + // create a VectorSchemaRoot root2 and load the recordBatch + VectorSchemaRoot root2 = VectorSchemaRoot.create(root1.getSchema(), allocator); + VectorLoader loader = new VectorLoader(root2); + loader.load(recordBatch); + +A new VectorSchemaRoot can be sliced from an existing root without copying +data: + +.. code-block:: Java + + // 0 indicates start index (inclusive) and 5 indicated length (exclusive). + VectorSchemaRoot newRoot = vectorSchemaRoot.slice(0, 5); + +Table +===== + +A `Table`_ is an immutable tabular data structure, very similar to VectorSchemaRoot, in that it is also built on ValueVectors and schemas. Unlike VectorSchemaRoot, Table is not designed for batch processing. Here is a version of the example above, showing how to create a Table, rather than a VectorSchemaRoot: + +.. code-block:: Java + + BitVector bitVector = new BitVector("boolean", allocator); + VarCharVector varCharVector = new VarCharVector("varchar", allocator); + bitVector.allocateNew(); + varCharVector.allocateNew(); + for (int i = 0; i < 10; i++) { + bitVector.setSafe(i, i % 2 == 0 ? 0 : 1); + varCharVector.setSafe(i, ("test" + i).getBytes(StandardCharsets.UTF_8)); + } + bitVector.setValueCount(10); + varCharVector.setValueCount(10); + + List vectors = Arrays.asList(bitVector, varCharVector); + Table table = new Table(vectors); + +See the :doc:`table` documentation for more information. + +.. _`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 9dac97dd9c..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.49.0 + 2.72.0 test diff --git a/flight/flight-core/src/main/java/module-info.java b/flight/flight-core/src/main/java/module-info.java index 28dbb732c4..669797ac93 100644 --- a/flight/flight-core/src/main/java/module-info.java +++ b/flight/flight-core/src/main/java/module-info.java @@ -20,6 +20,7 @@ exports org.apache.arrow.flight.auth; exports org.apache.arrow.flight.auth2; exports org.apache.arrow.flight.client; + exports org.apache.arrow.flight.grpc; exports org.apache.arrow.flight.impl; exports org.apache.arrow.flight.sql.impl; 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/FlightClient.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java index a15c3049aa..fd6e498d13 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java @@ -23,19 +23,13 @@ import io.grpc.ManagedChannel; import io.grpc.MethodDescriptor; import io.grpc.StatusRuntimeException; -import io.grpc.netty.GrpcSslContexts; import io.grpc.netty.NettyChannelBuilder; import io.grpc.stub.ClientCallStreamObserver; import io.grpc.stub.ClientCalls; import io.grpc.stub.ClientResponseObserver; import io.grpc.stub.StreamObserver; -import io.netty.channel.EventLoopGroup; -import io.netty.channel.ServerChannel; -import io.netty.handler.ssl.SslContextBuilder; -import io.netty.handler.ssl.util.InsecureTrustManagerFactory; import java.io.IOException; import java.io.InputStream; -import java.lang.reflect.InvocationTargetException; import java.net.URISyntaxException; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -45,7 +39,6 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.function.BooleanSupplier; -import javax.net.ssl.SSLException; import org.apache.arrow.flight.FlightProducer.StreamListener; import org.apache.arrow.flight.auth.BasicClientAuthHandler; import org.apache.arrow.flight.auth.ClientAuthHandler; @@ -57,6 +50,7 @@ import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; import org.apache.arrow.flight.grpc.ClientInterceptorAdapter; import org.apache.arrow.flight.grpc.CredentialCallOption; +import org.apache.arrow.flight.grpc.NettyClientBuilder; import org.apache.arrow.flight.grpc.StatusUtils; import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.impl.Flight.Empty; @@ -72,11 +66,6 @@ /** Client for Flight services. */ public class FlightClient implements AutoCloseable { private static final int PENDING_REQUESTS = 5; - /** - * The maximum number of trace events to keep on the gRPC Channel. This value disables channel - * tracing. - */ - private static final int MAX_CHANNEL_TRACE_EVENTS = 0; private final BufferAllocator allocator; private final ManagedChannel channel; @@ -96,11 +85,12 @@ public class FlightClient implements AutoCloseable { List middleware) { this.allocator = incomingAllocator.newChildAllocator("flight-client", 0, Long.MAX_VALUE); this.channel = channel; - this.middleware = middleware; + // We need a mutable copy (shared between this class and ClientInterceptorAdapter) + this.middleware = new ArrayList<>(middleware); final ClientInterceptor[] interceptors; interceptors = - new ClientInterceptor[] {authInterceptor, new ClientInterceptorAdapter(middleware)}; + new ClientInterceptor[] {authInterceptor, new ClientInterceptorAdapter(this.middleware)}; // Create a channel with interceptors pre-applied for DoGet and DoPut Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptors); @@ -771,176 +761,71 @@ public static Builder builder(BufferAllocator allocator, Location location) { /** A builder for Flight clients. */ public static final class Builder { - private BufferAllocator allocator; - private Location location; - private boolean forceTls = false; - private int maxInboundMessageSize = FlightServer.MAX_GRPC_MESSAGE_SIZE; - private InputStream trustedCertificates = null; - private InputStream clientCertificate = null; - private InputStream clientKey = null; - private String overrideHostname = null; - private List middleware = new ArrayList<>(); - private boolean verifyServer = true; - - private Builder() {} + private final NettyClientBuilder builder; + + private Builder() { + this.builder = new NettyClientBuilder(); + } private Builder(BufferAllocator allocator, Location location) { - this.allocator = Preconditions.checkNotNull(allocator); - this.location = Preconditions.checkNotNull(location); + this.builder = new NettyClientBuilder(allocator, location); } /** Force the client to connect over TLS. */ public Builder useTls() { - this.forceTls = true; + builder.useTls(); return this; } /** Override the hostname checked for TLS. Use with caution in production. */ public Builder overrideHostname(final String hostname) { - this.overrideHostname = hostname; + builder.overrideHostname(hostname); return this; } /** Set the maximum inbound message size. */ public Builder maxInboundMessageSize(int maxSize) { - Preconditions.checkArgument(maxSize > 0); - this.maxInboundMessageSize = maxSize; + builder.maxInboundMessageSize(maxSize); return this; } /** Set the trusted TLS certificates. */ public Builder trustedCertificates(final InputStream stream) { - this.trustedCertificates = Preconditions.checkNotNull(stream); + builder.trustedCertificates(stream); return this; } /** Set the trusted TLS certificates. */ public Builder clientCertificate( final InputStream clientCertificate, final InputStream clientKey) { - Preconditions.checkNotNull(clientKey); - this.clientCertificate = Preconditions.checkNotNull(clientCertificate); - this.clientKey = Preconditions.checkNotNull(clientKey); + builder.clientCertificate(clientCertificate, clientKey); return this; } public Builder allocator(BufferAllocator allocator) { - this.allocator = Preconditions.checkNotNull(allocator); + builder.allocator(allocator); return this; } public Builder location(Location location) { - this.location = Preconditions.checkNotNull(location); + builder.location(location); return this; } public Builder intercept(FlightClientMiddleware.Factory factory) { - middleware.add(factory); + builder.intercept(factory); return this; } public Builder verifyServer(boolean verifyServer) { - this.verifyServer = verifyServer; + builder.verifyServer(verifyServer); return this; } /** Create the client from this builder. */ public FlightClient build() { - final NettyChannelBuilder builder; - - switch (location.getUri().getScheme()) { - case LocationSchemes.GRPC: - case LocationSchemes.GRPC_INSECURE: - case LocationSchemes.GRPC_TLS: - { - builder = NettyChannelBuilder.forAddress(location.toSocketAddress()); - break; - } - case LocationSchemes.GRPC_DOMAIN_SOCKET: - { - // The implementation is platform-specific, so we have to find the classes at runtime - builder = NettyChannelBuilder.forAddress(location.toSocketAddress()); - try { - try { - // Linux - builder.channelType( - Class.forName("io.netty.channel.epoll.EpollDomainSocketChannel") - .asSubclass(ServerChannel.class)); - final EventLoopGroup elg = - Class.forName("io.netty.channel.epoll.EpollEventLoopGroup") - .asSubclass(EventLoopGroup.class) - .getDeclaredConstructor() - .newInstance(); - builder.eventLoopGroup(elg); - } catch (ClassNotFoundException e) { - // BSD - builder.channelType( - Class.forName("io.netty.channel.kqueue.KQueueDomainSocketChannel") - .asSubclass(ServerChannel.class)); - final EventLoopGroup elg = - Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup") - .asSubclass(EventLoopGroup.class) - .getDeclaredConstructor() - .newInstance(); - builder.eventLoopGroup(elg); - } - } catch (ClassNotFoundException - | InstantiationException - | IllegalAccessException - | NoSuchMethodException - | InvocationTargetException e) { - throw new UnsupportedOperationException( - "Could not find suitable Netty native transport implementation for domain socket address."); - } - break; - } - default: - throw new IllegalArgumentException( - "Scheme is not supported: " + location.getUri().getScheme()); - } - - if (this.forceTls || LocationSchemes.GRPC_TLS.equals(location.getUri().getScheme())) { - builder.useTransportSecurity(); - - final boolean hasTrustedCerts = this.trustedCertificates != null; - final boolean hasKeyCertPair = this.clientCertificate != null && this.clientKey != null; - if (!this.verifyServer && (hasTrustedCerts || hasKeyCertPair)) { - throw new IllegalArgumentException( - "FlightClient has been configured to disable server verification, " - + "but certificate options have been specified."); - } - - final SslContextBuilder sslContextBuilder = GrpcSslContexts.forClient(); - - if (!this.verifyServer) { - sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); - } else if (this.trustedCertificates != null - || this.clientCertificate != null - || this.clientKey != null) { - if (this.trustedCertificates != null) { - sslContextBuilder.trustManager(this.trustedCertificates); - } - if (this.clientCertificate != null && this.clientKey != null) { - sslContextBuilder.keyManager(this.clientCertificate, this.clientKey); - } - } - try { - builder.sslContext(sslContextBuilder.build()); - } catch (SSLException e) { - throw new RuntimeException(e); - } - - if (this.overrideHostname != null) { - builder.overrideAuthority(this.overrideHostname); - } - } else { - builder.usePlaintext(); - } - - builder - .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS) - .maxInboundMessageSize(maxInboundMessageSize) - .maxInboundMetadataSize(maxInboundMessageSize); - return new FlightClient(allocator, builder.build(), middleware); + final NettyChannelBuilder channelBuilder = builder.build(); + return new FlightClient(builder.allocator(), channelBuilder.build(), builder.middleware()); } } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightGrpcUtils.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightGrpcUtils.java index 13e4f2f215..df5e29741b 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightGrpcUtils.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightGrpcUtils.java @@ -23,6 +23,7 @@ import io.grpc.ManagedChannel; import io.grpc.MethodDescriptor; import java.util.Collections; +import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; import org.apache.arrow.flight.auth.ServerAuthHandler; @@ -151,6 +152,19 @@ public static FlightClient createFlightClient( return new FlightClient(incomingAllocator, channel, Collections.emptyList()); } + /** + * Creates a Flight client. + * + * @param incomingAllocator Memory allocator + * @param channel provides a connection to a gRPC server. + */ + public static FlightClient createFlightClient( + BufferAllocator incomingAllocator, + ManagedChannel channel, + List middleware) { + return new FlightClient(incomingAllocator, channel, middleware); + } + /** * Creates a Flight client. * 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/auth/ClientAuthWrapper.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ClientAuthWrapper.java index dd62f7756e..19d4f77a03 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ClientAuthWrapper.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ClientAuthWrapper.java @@ -76,9 +76,7 @@ public AuthObserver() { @Override public void onNext(HandshakeResponse value) { ByteString payload = value.getPayload(); - if (payload != null) { - messages.add(payload.toByteArray()); - } + messages.add(payload.toByteArray()); } private Iterator iter = diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ServerAuthWrapper.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ServerAuthWrapper.java index 879a93a73b..fba9f9baba 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ServerAuthWrapper.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/auth/ServerAuthWrapper.java @@ -84,9 +84,7 @@ public AuthObserver(StreamObserver responseObserver) { @Override public void onNext(HandshakeRequest value) { ByteString payload = value.getPayload(); - if (payload != null) { - messages.add(payload.toByteArray()); - } + messages.add(payload.toByteArray()); } private Iterator iter = 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/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java new file mode 100644 index 0000000000..42cdaac016 --- /dev/null +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/NettyClientBuilder.java @@ -0,0 +1,232 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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 io.grpc.ManagedChannel; +import io.grpc.netty.GrpcSslContexts; +import io.grpc.netty.NettyChannelBuilder; +import io.netty.channel.EventLoopGroup; +import io.netty.channel.ServerChannel; +import io.netty.handler.ssl.SslContextBuilder; +import io.netty.handler.ssl.util.InsecureTrustManagerFactory; +import java.io.InputStream; +import java.lang.reflect.InvocationTargetException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.net.ssl.SSLException; +import org.apache.arrow.flight.FlightClientMiddleware; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.LocationSchemes; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.Preconditions; + +/** + * A wrapper around gRPC's Netty builder. + * + *

It is recommended to use the Netty channel builder directly with {@link + * org.apache.arrow.flight.FlightGrpcUtils#createFlightClient(BufferAllocator, ManagedChannel)}. + * However, this class provides an adapter that implements the existing Flight-specific builder + * interface but allows usage of the Netty builder as well. + */ +public class NettyClientBuilder { + /** + * The maximum number of trace events to keep on the gRPC Channel. This value disables channel + * tracing. + */ + private static final int MAX_CHANNEL_TRACE_EVENTS = 0; + + protected BufferAllocator allocator; + protected Location location; + protected boolean forceTls = false; + protected int maxInboundMessageSize = Integer.MAX_VALUE; + protected InputStream trustedCertificates = null; + protected InputStream clientCertificate = null; + protected InputStream clientKey = null; + protected String overrideHostname = null; + protected List middleware = new ArrayList<>(); + protected boolean verifyServer = true; + + public NettyClientBuilder() {} + + public NettyClientBuilder(BufferAllocator allocator, Location location) { + this.allocator = Preconditions.checkNotNull(allocator); + this.location = Preconditions.checkNotNull(location); + } + + /** Force the client to connect over TLS. */ + public NettyClientBuilder useTls() { + this.forceTls = true; + return this; + } + + /** Override the hostname checked for TLS. Use with caution in production. */ + public NettyClientBuilder overrideHostname(final String hostname) { + this.overrideHostname = hostname; + return this; + } + + /** Set the maximum inbound message size. */ + public NettyClientBuilder maxInboundMessageSize(int maxSize) { + Preconditions.checkArgument(maxSize > 0); + this.maxInboundMessageSize = maxSize; + return this; + } + + /** Set the trusted TLS certificates. */ + public NettyClientBuilder trustedCertificates(final InputStream stream) { + this.trustedCertificates = Preconditions.checkNotNull(stream); + return this; + } + + /** Set the trusted TLS certificates. */ + public NettyClientBuilder clientCertificate( + final InputStream clientCertificate, final InputStream clientKey) { + Preconditions.checkNotNull(clientKey); + this.clientCertificate = Preconditions.checkNotNull(clientCertificate); + this.clientKey = Preconditions.checkNotNull(clientKey); + return this; + } + + public BufferAllocator allocator() { + return allocator; + } + + public NettyClientBuilder allocator(BufferAllocator allocator) { + this.allocator = Preconditions.checkNotNull(allocator); + return this; + } + + public NettyClientBuilder location(Location location) { + this.location = Preconditions.checkNotNull(location); + return this; + } + + public List middleware() { + return Collections.unmodifiableList(middleware); + } + + public NettyClientBuilder intercept(FlightClientMiddleware.Factory factory) { + middleware.add(factory); + return this; + } + + public NettyClientBuilder verifyServer(boolean verifyServer) { + this.verifyServer = verifyServer; + return this; + } + + /** Create the client from this builder. */ + public NettyChannelBuilder build() { + final NettyChannelBuilder builder; + + switch (location.getUri().getScheme()) { + case LocationSchemes.GRPC: + case LocationSchemes.GRPC_INSECURE: + case LocationSchemes.GRPC_TLS: + { + builder = NettyChannelBuilder.forAddress(location.toSocketAddress()); + break; + } + case LocationSchemes.GRPC_DOMAIN_SOCKET: + { + // The implementation is platform-specific, so we have to find the classes at runtime + builder = NettyChannelBuilder.forAddress(location.toSocketAddress()); + try { + try { + // Linux + builder.channelType( + Class.forName("io.netty.channel.epoll.EpollDomainSocketChannel") + .asSubclass(ServerChannel.class)); + final EventLoopGroup elg = + Class.forName("io.netty.channel.epoll.EpollEventLoopGroup") + .asSubclass(EventLoopGroup.class) + .getDeclaredConstructor() + .newInstance(); + builder.eventLoopGroup(elg); + } catch (ClassNotFoundException e) { + // BSD + builder.channelType( + Class.forName("io.netty.channel.kqueue.KQueueDomainSocketChannel") + .asSubclass(ServerChannel.class)); + final EventLoopGroup elg = + Class.forName("io.netty.channel.kqueue.KQueueEventLoopGroup") + .asSubclass(EventLoopGroup.class) + .getDeclaredConstructor() + .newInstance(); + builder.eventLoopGroup(elg); + } + } catch (ClassNotFoundException + | InstantiationException + | IllegalAccessException + | NoSuchMethodException + | InvocationTargetException e) { + throw new UnsupportedOperationException( + "Could not find suitable Netty native transport implementation for domain socket address."); + } + break; + } + default: + throw new IllegalArgumentException( + "Scheme is not supported: " + location.getUri().getScheme()); + } + + if (this.forceTls || LocationSchemes.GRPC_TLS.equals(location.getUri().getScheme())) { + builder.useTransportSecurity(); + + final boolean hasTrustedCerts = this.trustedCertificates != null; + final boolean hasKeyCertPair = this.clientCertificate != null && this.clientKey != null; + if (!this.verifyServer && (hasTrustedCerts || hasKeyCertPair)) { + throw new IllegalArgumentException( + "FlightClient has been configured to disable server verification, " + + "but certificate options have been specified."); + } + + final SslContextBuilder sslContextBuilder = GrpcSslContexts.forClient(); + + if (!this.verifyServer) { + sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); + } else if (this.trustedCertificates != null + || this.clientCertificate != null + || this.clientKey != null) { + if (this.trustedCertificates != null) { + sslContextBuilder.trustManager(this.trustedCertificates); + } + if (this.clientCertificate != null && this.clientKey != null) { + sslContextBuilder.keyManager(this.clientCertificate, this.clientKey); + } + } + try { + builder.sslContext(sslContextBuilder.build()); + } catch (SSLException e) { + throw new RuntimeException(e); + } + + if (this.overrideHostname != null) { + builder.overrideAuthority(this.overrideHostname); + } + } else { + builder.usePlaintext(); + } + + builder + .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS) + .maxInboundMessageSize(maxInboundMessageSize) + .maxInboundMetadataSize(maxInboundMessageSize); + return builder; + } +} diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestBasicOperation.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestBasicOperation.java index 5e818e6f5d..039bc14641 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestBasicOperation.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestBasicOperation.java @@ -486,8 +486,7 @@ public void testProtobufSchemaCompatibility() throws Exception { // Should have no body buffers assertFalse(message.getBufs().iterator().hasNext()); final Flight.FlightData protobufData = - arrowMessageToProtobuf(marshaller, message) - .toBuilder() + arrowMessageToProtobuf(marshaller, message).toBuilder() .setDataBody(ByteString.EMPTY) .build(); assertEquals(0, protobufData.getDataBody().size()); 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/TestLargeMessage.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestLargeMessage.java index 9362e1f552..e5c5d707c2 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestLargeMessage.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestLargeMessage.java @@ -134,7 +134,8 @@ public Runnable acceptPut( CallContext context, FlightStream flightStream, StreamListener ackStream) { return () -> { try (VectorSchemaRoot root = flightStream.getRoot()) { - while (flightStream.next()) {; + while (flightStream.next()) { + ; } } }; 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 e43bcd0571..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 @@ -68,7 +68,9 @@ under the License. + + org.apache.maven.plugins maven-shade-plugin @@ -87,9 +89,27 @@ under the License. **/module-info.class + + *:* + + LICENSE.txt + NOTICE.txt + META-INF/*LICENSE* + META-INF/*NOTICE* + META-INF/license/* + + - + + + META-INF/LICENSE.txt + src/shade/LICENSE.txt + + + META-INF/NOTICE.txt + src/shade/NOTICE.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 new file mode 100644 index 0000000000..3367c36f79 --- /dev/null +++ b/flight/flight-integration-tests/src/shade/LICENSE.txt @@ -0,0 +1,1127 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-------------------------------------------------------------------------------- + +This binary artifact contains Jackson 2.18.3. + +Home page: https://github.com/FasterXML/jackson +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons Codec 1.18.0. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-codec/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Google Flatbuffers 25.2.10. + +Home page: https://flatbuffers.dev/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Netty 4.1.119.Final. + +Copyright: Copyright 2014 The Netty Project +Home page: https://netty.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains gRPC 1.71.0. + +Copyright: Copyright 2014 The gRPC Authors +Home page: https://grpc.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Animal Sniffer Annotations 1.24. + +Copyright: Copyright (c) 2009 codehaus.org. +Home page: https://www.mojohaus.org/animal-sniffer/animal-sniffer-annotations/ +License: MIT (https://github.com/mojohaus/animal-sniffer/blob/animal-sniffer-1.24/LICENSE) + +-------------------------------------------------------------------------------- + +This binary artifact contains Perfmark 0.27.0. + +Copyright: Copyright 2019 Google LLC +Home page: https://github.com/perfmark/perfmark +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Gson 2.11.0. + +Copyright: Copyright 2008 Google Inc. +Home page: https://github.com/google/gson +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +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.2. + +Copyright: Copyright 2008 Google Inc. All rights reserved. +Home page: https://protobuf.dev/ +License: https://github.com/protocolbuffers/protobuf/blob/v4.30.1/LICENSE (BSD) +License text: + +| Copyright 2008 Google Inc. All rights reserved. +| +| Redistribution and use in source and binary forms, with or without +| modification, are permitted provided that the following conditions are +| met: +| +| * Redistributions of source code must retain the above copyright +| notice, this list of conditions and the following disclaimer. +| * Redistributions in binary form must reproduce the above +| copyright notice, this list of conditions and the following disclaimer +| in the documentation and/or other materials provided with the +| distribution. +| * Neither the name of Google Inc. nor the names of its +| contributors may be used to endorse or promote products derived from +| this software without specific prior written permission. +| +| THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +| "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +| LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +| A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +| OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +| SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +| LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +| DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +| THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +| (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +| OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +| +| Code generated by the Protocol Buffer compiler is owned by the owner +| of the input file used when generating it. This code is not +| standalone and requires a support library to be linked with it. This +| support library is itself covered by the above license. + +-------------------------------------------------------------------------------- + +This binary artifact contains Javax Annotation 1.3.2. + +Home page: https://github.com/javaee/javax.annotation +License: CDDL 1.1 +License text: + +| COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL) Version 1.1 +| +| 1. Definitions. +| +| 1.1. "Contributor" means each individual or entity that creates or +| contributes to the creation of Modifications. +| +| 1.2. "Contributor Version" means the combination of the Original +| Software, prior Modifications used by a Contributor (if any), and +| the Modifications made by that particular Contributor. +| +| 1.3. "Covered Software" means (a) the Original Software, or (b) +| Modifications, or (c) the combination of files containing Original +| Software with files containing Modifications, in each case including +| portions thereof. +| +| 1.4. "Executable" means the Covered Software in any form other than +| Source Code. +| +| 1.5. "Initial Developer" means the individual or entity that first +| makes Original Software available under this License. +| +| 1.6. "Larger Work" means a work which combines Covered Software or +| portions thereof with code not governed by the terms of this License. +| +| 1.7. "License" means this document. +| +| 1.8. "Licensable" means having the right to grant, to the maximum +| extent possible, whether at the time of the initial grant or +| subsequently acquired, any and all of the rights conveyed herein. +| +| 1.9. "Modifications" means the Source Code and Executable form of +| any of the following: +| +| A. Any file that results from an addition to, deletion from or +| modification of the contents of a file containing Original Software +| or previous Modifications; +| +| B. Any new file that contains any part of the Original Software or +| previous Modification; or +| +| C. Any new file that is contributed or otherwise made available +| under the terms of this License. +| +| 1.10. "Original Software" means the Source Code and Executable form +| of computer software code that is originally released under this +| License. +| +| 1.11. "Patent Claims" means any patent claim(s), now owned or +| hereafter acquired, including without limitation, method, process, +| and apparatus claims, in any patent Licensable by grantor. +| +| 1.12. "Source Code" means (a) the common form of computer software +| code in which modifications are made and (b) associated +| documentation included in or with such code. +| +| 1.13. "You" (or "Your") means an individual or a legal entity +| exercising rights under, and complying with all of the terms of, +| this License. For legal entities, "You" includes any entity which +| controls, is controlled by, or is under common control with You. For +| purposes of this definition, "control" means (a) the power, direct +| or indirect, to cause the direction or management of such entity, +| whether by contract or otherwise, or (b) ownership of more than +| fifty percent (50%) of the outstanding shares or beneficial +| ownership of such entity. +| +| 2. License Grants. +| +| 2.1. The Initial Developer Grant. +| +| Conditioned upon Your compliance with Section 3.1 below and subject +| to third party intellectual property claims, the Initial Developer +| hereby grants You a world-wide, royalty-free, non-exclusive license: +| +| (a) under intellectual property rights (other than patent or +| trademark) Licensable by Initial Developer, to use, reproduce, +| modify, display, perform, sublicense and distribute the Original +| Software (or portions thereof), with or without Modifications, +| and/or as part of a Larger Work; and +| +| (b) under Patent Claims infringed by the making, using or selling of +| Original Software, to make, have made, use, practice, sell, and +| offer for sale, and/or otherwise dispose of the Original Software +| (or portions thereof). +| +| (c) The licenses granted in Sections 2.1(a) and (b) are effective on +| the date Initial Developer first distributes or otherwise makes the +| Original Software available to a third party under the terms of this +| License. +| +| (d) Notwithstanding Section 2.1(b) above, no patent license is +| granted: (1) for code that You delete from the Original Software, or +| (2) for infringements caused by: (i) the modification of the +| Original Software, or (ii) the combination of the Original Software +| with other software or devices. +| +| 2.2. Contributor Grant. +| +| Conditioned upon Your compliance with Section 3.1 below and subject +| to third party intellectual property claims, each Contributor hereby +| grants You a world-wide, royalty-free, non-exclusive license: +| +| (a) under intellectual property rights (other than patent or +| trademark) Licensable by Contributor to use, reproduce, modify, +| display, perform, sublicense and distribute the Modifications +| created by such Contributor (or portions thereof), either on an +| unmodified basis, with other Modifications, as Covered Software +| and/or as part of a Larger Work; and +| +| (b) under Patent Claims infringed by the making, using, or selling +| of Modifications made by that Contributor either alone and/or in +| combination with its Contributor Version (or portions of such +| combination), to make, use, sell, offer for sale, have made, and/or +| otherwise dispose of: (1) Modifications made by that Contributor (or +| portions thereof); and (2) the combination of Modifications made by +| that Contributor with its Contributor Version (or portions of such +| combination). +| +| (c) The licenses granted in Sections 2.2(a) and 2.2(b) are effective +| on the date Contributor first distributes or otherwise makes the +| Modifications available to a third party. +| +| (d) Notwithstanding Section 2.2(b) above, no patent license is +| granted: (1) for any code that Contributor has deleted from the +| Contributor Version; (2) for infringements caused by: (i) third +| party modifications of Contributor Version, or (ii) the combination +| of Modifications made by that Contributor with other software +| (except as part of the Contributor Version) or other devices; or (3) +| under Patent Claims infringed by Covered Software in the absence of +| Modifications made by that Contributor. +| +| 3. Distribution Obligations. +| +| 3.1. Availability of Source Code. +| +| Any Covered Software that You distribute or otherwise make available +| in Executable form must also be made available in Source Code form +| and that Source Code form must be distributed only under the terms +| of this License. You must include a copy of this License with every +| copy of the Source Code form of the Covered Software You distribute +| or otherwise make available. You must inform recipients of any such +| Covered Software in Executable form as to how they can obtain such +| Covered Software in Source Code form in a reasonable manner on or +| through a medium customarily used for software exchange. +| +| 3.2. Modifications. +| +| The Modifications that You create or to which You contribute are +| governed by the terms of this License. You represent that You +| believe Your Modifications are Your original creation(s) and/or You +| have sufficient rights to grant the rights conveyed by this License. +| +| 3.3. Required Notices. +| +| You must include a notice in each of Your Modifications that +| identifies You as the Contributor of the Modification. You may not +| remove or alter any copyright, patent or trademark notices contained +| within the Covered Software, or any notices of licensing or any +| descriptive text giving attribution to any Contributor or the +| Initial Developer. +| +| 3.4. Application of Additional Terms. +| +| You may not offer or impose any terms on any Covered Software in +| Source Code form that alters or restricts the applicable version of +| this License or the recipients' rights hereunder. You may choose to +| offer, and to charge a fee for, warranty, support, indemnity or +| liability obligations to one or more recipients of Covered Software. +| However, you may do so only on Your own behalf, and not on behalf of +| the Initial Developer or any Contributor. You must make it +| absolutely clear that any such warranty, support, indemnity or +| liability obligation is offered by You alone, and You hereby agree +| to indemnify the Initial Developer and every Contributor for any +| liability incurred by the Initial Developer or such Contributor as a +| result of warranty, support, indemnity or liability terms You offer. +| +| 3.5. Distribution of Executable Versions. +| +| You may distribute the Executable form of the Covered Software under +| the terms of this License or under the terms of a license of Your +| choice, which may contain terms different from this License, +| provided that You are in compliance with the terms of this License +| and that the license for the Executable form does not attempt to +| limit or alter the recipient's rights in the Source Code form from +| the rights set forth in this License. If You distribute the Covered +| Software in Executable form under a different license, You must make +| it absolutely clear that any terms which differ from this License +| are offered by You alone, not by the Initial Developer or +| Contributor. You hereby agree to indemnify the Initial Developer and +| every Contributor for any liability incurred by the Initial +| Developer or such Contributor as a result of any such terms You offer. +| +| 3.6. Larger Works. +| +| You may create a Larger Work by combining Covered Software with +| other code not governed by the terms of this License and distribute +| the Larger Work as a single product. In such a case, You must make +| sure the requirements of this License are fulfilled for the Covered +| Software. +| +| 4. Versions of the License. +| +| 4.1. New Versions. +| +| Oracle is the initial license steward and may publish revised and/or +| new versions of this License from time to time. Each version will be +| given a distinguishing version number. Except as provided in Section +| 4.3, no one other than the license steward has the right to modify +| this License. +| +| 4.2. Effect of New Versions. +| +| You may always continue to use, distribute or otherwise make the +| Covered Software available under the terms of the version of the +| License under which You originally received the Covered Software. If +| the Initial Developer includes a notice in the Original Software +| prohibiting it from being distributed or otherwise made available +| under any subsequent version of the License, You must distribute and +| make the Covered Software available under the terms of the version +| of the License under which You originally received the Covered +| Software. Otherwise, You may also choose to use, distribute or +| otherwise make the Covered Software available under the terms of any +| subsequent version of the License published by the license steward. +| +| 4.3. Modified Versions. +| +| When You are an Initial Developer and You want to create a new +| license for Your Original Software, You may create and use a +| modified version of this License if You: (a) rename the license and +| remove any references to the name of the license steward (except to +| note that the license differs from this License); and (b) otherwise +| make it clear that the license contains terms which differ from this +| License. +| +| 5. DISCLAIMER OF WARRANTY. +| +| COVERED SOFTWARE IS PROVIDED UNDER THIS LICENSE ON AN "AS IS" BASIS, +| WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, +| INCLUDING, WITHOUT LIMITATION, WARRANTIES THAT THE COVERED SOFTWARE +| IS FREE OF DEFECTS, MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE OR +| NON-INFRINGING. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF +| THE COVERED SOFTWARE IS WITH YOU. SHOULD ANY COVERED SOFTWARE PROVE +| DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL DEVELOPER OR ANY +| OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY SERVICING, +| REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN +| ESSENTIAL PART OF THIS LICENSE. NO USE OF ANY COVERED SOFTWARE IS +| AUTHORIZED HEREUNDER EXCEPT UNDER THIS DISCLAIMER. +| +| 6. TERMINATION. +| +| 6.1. This License and the rights granted hereunder will terminate +| automatically if You fail to comply with terms herein and fail to +| cure such breach within 30 days of becoming aware of the breach. +| Provisions which, by their nature, must remain in effect beyond the +| termination of this License shall survive. +| +| 6.2. If You assert a patent infringement claim (excluding +| declaratory judgment actions) against Initial Developer or a +| Contributor (the Initial Developer or Contributor against whom You +| assert such claim is referred to as "Participant") alleging that the +| Participant Software (meaning the Contributor Version where the +| Participant is a Contributor or the Original Software where the +| Participant is the Initial Developer) directly or indirectly +| infringes any patent, then any and all rights granted directly or +| indirectly to You by such Participant, the Initial Developer (if the +| Initial Developer is not the Participant) and all Contributors under +| Sections 2.1 and/or 2.2 of this License shall, upon 60 days notice +| from Participant terminate prospectively and automatically at the +| expiration of such 60 day notice period, unless if within such 60 +| day period You withdraw Your claim with respect to the Participant +| Software against such Participant either unilaterally or pursuant to +| a written agreement with Participant. +| +| 6.3. If You assert a patent infringement claim against Participant +| alleging that the Participant Software directly or indirectly +| infringes any patent where such claim is resolved (such as by +| license or settlement) prior to the initiation of patent +| infringement litigation, then the reasonable value of the licenses +| granted by such Participant under Sections 2.1 or 2.2 shall be taken +| into account in determining the amount or value of any payment or +| license. +| +| 6.4. In the event of termination under Sections 6.1 or 6.2 above, +| all end user licenses that have been validly granted by You or any +| distributor hereunder prior to termination (excluding licenses +| granted to You by any distributor) shall survive termination. +| +| 7. LIMITATION OF LIABILITY. +| +| UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL THEORY, WHETHER TORT +| (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE, SHALL YOU, THE +| INITIAL DEVELOPER, ANY OTHER CONTRIBUTOR, OR ANY DISTRIBUTOR OF +| COVERED SOFTWARE, OR ANY SUPPLIER OF ANY OF SUCH PARTIES, BE LIABLE +| TO ANY PERSON FOR ANY INDIRECT, SPECIAL, INCIDENTAL, OR +| CONSEQUENTIAL DAMAGES OF ANY CHARACTER INCLUDING, WITHOUT +| LIMITATION, DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER +| FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR +| LOSSES, EVEN IF SUCH PARTY SHALL HAVE BEEN INFORMED OF THE +| POSSIBILITY OF SUCH DAMAGES. THIS LIMITATION OF LIABILITY SHALL NOT +| APPLY TO LIABILITY FOR DEATH OR PERSONAL INJURY RESULTING FROM SUCH +| PARTY'S NEGLIGENCE TO THE EXTENT APPLICABLE LAW PROHIBITS SUCH +| LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR +| LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION +| AND LIMITATION MAY NOT APPLY TO YOU. +| +| 8. U.S. GOVERNMENT END USERS. +| +| The Covered Software is a "commercial item," as that term is defined +| in 48 C.F.R. 2.101 (Oct. 1995), consisting of "commercial computer +| software" (as that term is defined at 48 C.F.R. § +| 252.227-7014(a)(1)) and "commercial computer software documentation" +| as such terms are used in 48 C.F.R. 12.212 (Sept. 1995). Consistent +| with 48 C.F.R. 12.212 and 48 C.F.R. 227.7202-1 through 227.7202-4 +| (June 1995), all U.S. Government End Users acquire Covered Software +| with only those rights set forth herein. This U.S. Government Rights +| clause is in lieu of, and supersedes, any other FAR, DFAR, or other +| clause or provision that addresses Government rights in computer +| software under this License. +| +| 9. MISCELLANEOUS. +| +| This License represents the complete agreement concerning subject +| matter hereof. If any provision of this License is held to be +| unenforceable, such provision shall be reformed only to the extent +| necessary to make it enforceable. This License shall be governed by +| the law of the jurisdiction specified in a notice contained within +| the Original Software (except to the extent applicable law, if any, +| provides otherwise), excluding such jurisdiction's conflict-of-law +| provisions. Any litigation relating to this License shall be subject +| to the jurisdiction of the courts located in the jurisdiction and +| venue specified in a notice contained within the Original Software, +| with the losing party responsible for costs, including, without +| limitation, court costs and reasonable attorneys' fees and expenses. +| The application of the United Nations Convention on Contracts for +| the International Sale of Goods is expressly excluded. Any law or +| regulation which provides that the language of a contract shall be +| construed against the drafter shall not apply to this License. You +| agree that You alone are responsible for compliance with the United +| States export administration regulations (and the export control +| laws and regulation of any other countries) when You use, distribute +| or otherwise make available any Covered Software. +| +| 10. RESPONSIBILITY FOR CLAIMS. +| +| As between Initial Developer and the Contributors, each party is +| responsible for claims and damages arising, directly or indirectly, +| out of its utilization of rights under this License and You agree to +| work with Initial Developer and Contributors to distribute such +| responsibility on an equitable basis. Nothing herein is intended or +| shall be deemed to constitute any admission of liability. +| +| ------------------------------------------------------------------------ +| +| NOTICE PURSUANT TO SECTION 9 OF THE COMMON DEVELOPMENT AND DISTRIBUTION +| LICENSE (CDDL) +| +| The code released under the CDDL shall be governed by the laws of the +| State of California (excluding conflict-of-law provisions). Any +| litigation relating to this License shall be subject to the jurisdiction +| of the Federal Courts of the Northern District of California and the +| state courts of the State of California, with venue lying in Santa Clara +| County, California. +| +| +| +| The GNU General Public License (GPL) Version 2, June 1991 +| +| Copyright (C) 1989, 1991 Free Software Foundation, Inc. +| 51 Franklin Street, Fifth Floor +| Boston, MA 02110-1335 +| USA +| +| Everyone is permitted to copy and distribute verbatim copies +| of this license document, but changing it is not allowed. +| +| Preamble +| +| The licenses for most software are designed to take away your freedom to +| share and change it. By contrast, the GNU General Public License is +| intended to guarantee your freedom to share and change free software--to +| make sure the software is free for all its users. This General Public +| License applies to most of the Free Software Foundation's software and +| to any other program whose authors commit to using it. (Some other Free +| Software Foundation software is covered by the GNU Library General +| Public License instead.) You can apply it to your programs, too. +| +| When we speak of free software, we are referring to freedom, not price. +| Our General Public Licenses are designed to make sure that you have the +| freedom to distribute copies of free software (and charge for this +| service if you wish), that you receive source code or can get it if you +| want it, that you can change the software or use pieces of it in new +| free programs; and that you know you can do these things. +| +| To protect your rights, we need to make restrictions that forbid anyone +| to deny you these rights or to ask you to surrender the rights. These +| restrictions translate to certain responsibilities for you if you +| distribute copies of the software, or if you modify it. +| +| For example, if you distribute copies of such a program, whether gratis +| or for a fee, you must give the recipients all the rights that you have. +| You must make sure that they, too, receive or can get the source code. +| And you must show them these terms so they know their rights. +| +| We protect your rights with two steps: (1) copyright the software, and +| (2) offer you this license which gives you legal permission to copy, +| distribute and/or modify the software. +| +| Also, for each author's protection and ours, we want to make certain +| that everyone understands that there is no warranty for this free +| software. If the software is modified by someone else and passed on, we +| want its recipients to know that what they have is not the original, so +| that any problems introduced by others will not reflect on the original +| authors' reputations. +| +| Finally, any free program is threatened constantly by software patents. +| We wish to avoid the danger that redistributors of a free program will +| individually obtain patent licenses, in effect making the program +| proprietary. To prevent this, we have made it clear that any patent must +| be licensed for everyone's free use or not licensed at all. +| +| The precise terms and conditions for copying, distribution and +| modification follow. +| +| TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION +| +| 0. This License applies to any program or other work which contains a +| notice placed by the copyright holder saying it may be distributed under +| the terms of this General Public License. The "Program", below, refers +| to any such program or work, and a "work based on the Program" means +| either the Program or any derivative work under copyright law: that is +| to say, a work containing the Program or a portion of it, either +| verbatim or with modifications and/or translated into another language. +| (Hereinafter, translation is included without limitation in the term +| "modification".) Each licensee is addressed as "you". +| +| Activities other than copying, distribution and modification are not +| covered by this License; they are outside its scope. The act of running +| the Program is not restricted, and the output from the Program is +| covered only if its contents constitute a work based on the Program +| (independent of having been made by running the Program). Whether that +| is true depends on what the Program does. +| +| 1. You may copy and distribute verbatim copies of the Program's source +| code as you receive it, in any medium, provided that you conspicuously +| and appropriately publish on each copy an appropriate copyright notice +| and disclaimer of warranty; keep intact all the notices that refer to +| this License and to the absence of any warranty; and give any other +| recipients of the Program a copy of this License along with the Program. +| +| You may charge a fee for the physical act of transferring a copy, and +| you may at your option offer warranty protection in exchange for a fee. +| +| 2. You may modify your copy or copies of the Program or any portion of +| it, thus forming a work based on the Program, and copy and distribute +| such modifications or work under the terms of Section 1 above, provided +| that you also meet all of these conditions: +| +| a) You must cause the modified files to carry prominent notices +| stating that you changed the files and the date of any change. +| +| b) You must cause any work that you distribute or publish, that in +| whole or in part contains or is derived from the Program or any part +| thereof, to be licensed as a whole at no charge to all third parties +| under the terms of this License. +| +| c) If the modified program normally reads commands interactively +| when run, you must cause it, when started running for such +| interactive use in the most ordinary way, to print or display an +| announcement including an appropriate copyright notice and a notice +| that there is no warranty (or else, saying that you provide a +| warranty) and that users may redistribute the program under these +| conditions, and telling the user how to view a copy of this License. +| (Exception: if the Program itself is interactive but does not +| normally print such an announcement, your work based on the Program +| is not required to print an announcement.) +| +| These requirements apply to the modified work as a whole. If +| identifiable sections of that work are not derived from the Program, and +| can be reasonably considered independent and separate works in +| themselves, then this License, and its terms, do not apply to those +| sections when you distribute them as separate works. But when you +| distribute the same sections as part of a whole which is a work based on +| the Program, the distribution of the whole must be on the terms of this +| License, whose permissions for other licensees extend to the entire +| whole, and thus to each and every part regardless of who wrote it. +| +| Thus, it is not the intent of this section to claim rights or contest +| your rights to work written entirely by you; rather, the intent is to +| exercise the right to control the distribution of derivative or +| collective works based on the Program. +| +| In addition, mere aggregation of another work not based on the Program +| with the Program (or with a work based on the Program) on a volume of a +| storage or distribution medium does not bring the other work under the +| scope of this License. +| +| 3. You may copy and distribute the Program (or a work based on it, +| under Section 2) in object code or executable form under the terms of +| Sections 1 and 2 above provided that you also do one of the following: +| +| a) Accompany it with the complete corresponding machine-readable +| source code, which must be distributed under the terms of Sections 1 +| and 2 above on a medium customarily used for software interchange; or, +| +| b) Accompany it with a written offer, valid for at least three +| years, to give any third party, for a charge no more than your cost +| of physically performing source distribution, a complete +| machine-readable copy of the corresponding source code, to be +| distributed under the terms of Sections 1 and 2 above on a medium +| customarily used for software interchange; or, +| +| c) Accompany it with the information you received as to the offer to +| distribute corresponding source code. (This alternative is allowed +| only for noncommercial distribution and only if you received the +| program in object code or executable form with such an offer, in +| accord with Subsection b above.) +| +| The source code for a work means the preferred form of the work for +| making modifications to it. For an executable work, complete source code +| means all the source code for all modules it contains, plus any +| associated interface definition files, plus the scripts used to control +| compilation and installation of the executable. However, as a special +| exception, the source code distributed need not include anything that is +| normally distributed (in either source or binary form) with the major +| components (compiler, kernel, and so on) of the operating system on +| which the executable runs, unless that component itself accompanies the +| executable. +| +| If distribution of executable or object code is made by offering access +| to copy from a designated place, then offering equivalent access to copy +| the source code from the same place counts as distribution of the source +| code, even though third parties are not compelled to copy the source +| along with the object code. +| +| 4. You may not copy, modify, sublicense, or distribute the Program +| except as expressly provided under this License. Any attempt otherwise +| to copy, modify, sublicense or distribute the Program is void, and will +| automatically terminate your rights under this License. However, parties +| who have received copies, or rights, from you under this License will +| not have their licenses terminated so long as such parties remain in +| full compliance. +| +| 5. You are not required to accept this License, since you have not +| signed it. However, nothing else grants you permission to modify or +| distribute the Program or its derivative works. These actions are +| prohibited by law if you do not accept this License. Therefore, by +| modifying or distributing the Program (or any work based on the +| Program), you indicate your acceptance of this License to do so, and all +| its terms and conditions for copying, distributing or modifying the +| Program or works based on it. +| +| 6. Each time you redistribute the Program (or any work based on the +| Program), the recipient automatically receives a license from the +| original licensor to copy, distribute or modify the Program subject to +| these terms and conditions. You may not impose any further restrictions +| on the recipients' exercise of the rights granted herein. You are not +| responsible for enforcing compliance by third parties to this License. +| +| 7. If, as a consequence of a court judgment or allegation of patent +| infringement or for any other reason (not limited to patent issues), +| conditions are imposed on you (whether by court order, agreement or +| otherwise) that contradict the conditions of this License, they do not +| excuse you from the conditions of this License. If you cannot distribute +| so as to satisfy simultaneously your obligations under this License and +| any other pertinent obligations, then as a consequence you may not +| distribute the Program at all. For example, if a patent license would +| not permit royalty-free redistribution of the Program by all those who +| receive copies directly or indirectly through you, then the only way you +| could satisfy both it and this License would be to refrain entirely from +| distribution of the Program. +| +| If any portion of this section is held invalid or unenforceable under +| any particular circumstance, the balance of the section is intended to +| apply and the section as a whole is intended to apply in other +| circumstances. +| +| It is not the purpose of this section to induce you to infringe any +| patents or other property right claims or to contest validity of any +| such claims; this section has the sole purpose of protecting the +| integrity of the free software distribution system, which is implemented +| by public license practices. Many people have made generous +| contributions to the wide range of software distributed through that +| system in reliance on consistent application of that system; it is up to +| the author/donor to decide if he or she is willing to distribute +| software through any other system and a licensee cannot impose that choice. +| +| This section is intended to make thoroughly clear what is believed to be +| a consequence of the rest of this License. +| +| 8. If the distribution and/or use of the Program is restricted in +| certain countries either by patents or by copyrighted interfaces, the +| original copyright holder who places the Program under this License may +| add an explicit geographical distribution limitation excluding those +| countries, so that distribution is permitted only in or among countries +| not thus excluded. In such case, this License incorporates the +| limitation as if written in the body of this License. +| +| 9. The Free Software Foundation may publish revised and/or new +| versions of the General Public License from time to time. Such new +| versions will be similar in spirit to the present version, but may +| differ in detail to address new problems or concerns. +| +| Each version is given a distinguishing version number. If the Program +| specifies a version number of this License which applies to it and "any +| later version", you have the option of following the terms and +| conditions either of that version or of any later version published by +| the Free Software Foundation. If the Program does not specify a version +| number of this License, you may choose any version ever published by the +| Free Software Foundation. +| +| 10. If you wish to incorporate parts of the Program into other free +| programs whose distribution conditions are different, write to the +| author to ask for permission. For software which is copyrighted by the +| Free Software Foundation, write to the Free Software Foundation; we +| sometimes make exceptions for this. Our decision will be guided by the +| two goals of preserving the free status of all derivatives of our free +| software and of promoting the sharing and reuse of software generally. +| +| NO WARRANTY +| +| 11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO +| WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. +| EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR +| OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, +| EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +| WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE +| ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH +| YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL +| NECESSARY SERVICING, REPAIR OR CORRECTION. +| +| 12. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN +| WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY +| AND/OR REDISTRIBUTE THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR +| DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL +| DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM +| (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED +| INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF +| THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR +| OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. +| +| END OF TERMS AND CONDITIONS +| +| How to Apply These Terms to Your New Programs +| +| If you develop a new program, and you want it to be of the greatest +| possible use to the public, the best way to achieve this is to make it +| free software which everyone can redistribute and change under these terms. +| +| To do so, attach the following notices to the program. It is safest to +| attach them to the start of each source file to most effectively convey +| the exclusion of warranty; and each file should have at least the +| "copyright" line and a pointer to where the full notice is found. +| +| One line to give the program's name and a brief idea of what it does. +| Copyright (C) +| +| This program is free software; you can redistribute it and/or modify +| it under the terms of the GNU General Public License as published by +| the Free Software Foundation; either version 2 of the License, or +| (at your option) any later version. +| +| This program is distributed in the hope that it will be useful, but +| WITHOUT ANY WARRANTY; without even the implied warranty of +| MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +| General Public License for more details. +| +| You should have received a copy of the GNU General Public License +| along with this program; if not, write to the Free Software +| Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335 USA +| +| Also add information on how to contact you by electronic and paper mail. +| +| If the program is interactive, make it output a short notice like this +| when it starts in an interactive mode: +| +| Gnomovision version 69, Copyright (C) year name of author +| Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type +| `show w'. This is free software, and you are welcome to redistribute +| it under certain conditions; type `show c' for details. +| +| The hypothetical commands `show w' and `show c' should show the +| appropriate parts of the General Public License. Of course, the commands +| you use may be called something other than `show w' and `show c'; they +| could even be mouse-clicks or menu items--whatever suits your program. +| +| You should also get your employer (if you work as a programmer) or your +| school, if any, to sign a "copyright disclaimer" for the program, if +| necessary. Here is a sample; alter the names: +| +| Yoyodyne, Inc., hereby disclaims all copyright interest in the +| program `Gnomovision' (which makes passes at compilers) written by +| James Hacker. +| +| signature of Ty Coon, 1 April 1989 +| Ty Coon, President of Vice +| +| This General Public License does not permit incorporating your program +| into proprietary programs. If your program is a subroutine library, you +| may consider it more useful to permit linking proprietary applications +| with the library. If this is what you want to do, use the GNU Library +| General Public License instead of this License. +| +| # +| +| Certain source files distributed by Oracle America, Inc. and/or its +| affiliates are subject to the following clarification and special +| exception to the GPLv2, based on the GNU Project exception for its +| Classpath libraries, known as the GNU Classpath Exception, but only +| where Oracle has expressly included in the particular source file's +| header the words "Oracle designates this particular file as subject to +| the "Classpath" exception as provided by Oracle in the LICENSE file +| that accompanied this code." +| +| You should also note that Oracle includes multiple, independent +| programs in this software package. Some of those programs are provided +| under licenses deemed incompatible with the GPLv2 by the Free Software +| Foundation and others. For example, the package includes programs +| licensed under the Apache License, Version 2.0. Such programs are +| licensed to you under their original licenses. +| +| Oracle facilitates your further distribution of this package by adding +| the Classpath Exception to the necessary parts of its GPLv2 code, which +| permits you to use that code in combination with other independent +| modules not licensed under the GPLv2. However, note that this would +| not permit you to commingle code under an incompatible license with +| Oracle's GPLv2 licensed code by, for example, cutting and pasting such +| code into a file also containing Oracle's GPLv2 licensed code and then +| distributing the result. Additionally, if you were to remove the +| Classpath Exception from any of the files to which it applies and +| distribute the result, you would likely be required to license some or +| all of the other code in that distribution under the GPLv2 as well, and +| since the GPLv2 is incompatible with the license terms of some items +| included in the distribution by Oracle, removing the Classpath +| Exception could therefore effectively compromise your ability to +| further distribute the package. +| +| Proceed with caution and we recommend that you obtain the advice of a +| lawyer skilled in open source matters before removing the Classpath +| Exception or making modifications to this package which may +| subsequently be redistributed and/or involve the use of third party +| software. +| +| CLASSPATH EXCEPTION +| Linking this library statically or dynamically with other modules is +| making a combined work based on this library. Thus, the terms and +| conditions of the GNU General Public License version 2 cover the whole +| combination. +| +| As a special exception, the copyright holders of this library give you +| permission to link this library with independent modules to produce an +| executable, regardless of the license terms of these independent +| modules, and to copy and distribute the resulting executable under +| terms of your choice, provided that you also meet, for each linked +| independent module, the terms and conditions of the license of that +| module. An independent module is a module which is not derived from or +| based on this library. If you modify this library, you may extend this +| exception to your version of the library, but you are not obligated to +| do so. If you do not wish to do so, delete this exception statement +| from your version. + +-------------------------------------------------------------------------------- + +This binary artifact contains Google j2objc 3.0.0. + +Home page: http://j2objc.org/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons CLI 1.9.0. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-cli/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Slf4j 2.0.17. + +Copyright: Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +Home page: http://www.slf4j.org/ +License: MIT +License text: + +| Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +| All rights reserved. +| +| Permission is hereby granted, free of charge, to any person obtaining +| a copy of this software and associated documentation files (the +| "Software"), to deal in the Software without restriction, including +| without limitation the rights to use, copy, modify, merge, publish, +| distribute, sublicense, and/or sell copies of the Software, and to +| permit persons to whom the Software is furnished to do so, subject to +| the following conditions: +| +| The above copyright notice and this permission notice shall be +| included in all copies or substantial portions of the Software. +| +| THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +| EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +| MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +| NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +| LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +| OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +| WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/flight/flight-integration-tests/src/shade/NOTICE.txt b/flight/flight-integration-tests/src/shade/NOTICE.txt new file mode 100644 index 0000000000..668cfe684b --- /dev/null +++ b/flight/flight-integration-tests/src/shade/NOTICE.txt @@ -0,0 +1,340 @@ +Apache Arrow Java +Copyright 2016-2025 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +--------------------------------------------------------- + +This product includes Netty 4.1.119.Final, with the following in its NOTICE: + +| The Netty Project +| ================= +| +| Please visit the Netty web site for more information: +| +| * https://netty.io/ +| +| Copyright 2014 The Netty Project +| +| The Netty Project licenses this file to you under the Apache License, +| version 2.0 (the "License"); you may not use this file except in compliance +| with the License. You may obtain a copy of the License at: +| +| https://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +| WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +| License for the specific language governing permissions and limitations +| under the License. +| +| Also, please refer to each LICENSE..txt file, which is located in +| the 'license' directory of the distribution file, for the license terms of the +| components that this product depends on. +| +| ------------------------------------------------------------------------------- +| This product contains the extensions to Java Collections Framework which has +| been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene: +| +| * LICENSE: +| * license/LICENSE.jsr166y.txt (Public Domain) +| * HOMEPAGE: +| * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/ +| * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/ +| +| This product contains a modified version of Robert Harder's Public Domain +| Base64 Encoder and Decoder, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.base64.txt (Public Domain) +| * HOMEPAGE: +| * http://iharder.sourceforge.net/current/java/base64/ +| +| This product contains a modified portion of 'Webbit', an event based +| WebSocket and HTTP server, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.webbit.txt (BSD License) +| * HOMEPAGE: +| * https://github.com/joewalnes/webbit +| +| This product contains a modified portion of 'SLF4J', a simple logging +| facade for Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.slf4j.txt (MIT License) +| * HOMEPAGE: +| * https://www.slf4j.org/ +| +| This product contains a modified portion of 'Apache Harmony', an open source +| Java SE, which can be obtained at: +| +| * NOTICE: +| * license/NOTICE.harmony.txt +| * LICENSE: +| * license/LICENSE.harmony.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://archive.apache.org/dist/harmony/ +| +| This product contains a modified portion of 'jbzip2', a Java bzip2 compression +| and decompression library written by Matthew J. Francis. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jbzip2.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jbzip2/ +| +| This product contains a modified portion of 'libdivsufsort', a C API library to construct +| the suffix array and the Burrows-Wheeler transformed string for any input string of +| a constant-size alphabet written by Yuta Mori. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.libdivsufsort.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/y-256/libdivsufsort +| +| This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM, +| which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jctools.txt (ASL2 License) +| * HOMEPAGE: +| * https://github.com/JCTools/JCTools +| +| This product optionally depends on 'JZlib', a re-implementation of zlib in +| pure Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jzlib.txt (BSD style License) +| * HOMEPAGE: +| * http://www.jcraft.com/jzlib/ +| +| This product optionally depends on 'Compress-LZF', a Java library for encoding and +| decoding data in LZF format, written by Tatu Saloranta. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.compress-lzf.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/ning/compress +| +| This product optionally depends on 'lz4', a LZ4 Java compression +| and decompression library written by Adrien Grand. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lz4.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jpountz/lz4-java +| +| This product optionally depends on 'lzma-java', a LZMA Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lzma-java.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jponge/lzma-java +| +| This product optionally depends on 'zstd-jni', a zstd-jni Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.zstd-jni.txt (BSD) +| * HOMEPAGE: +| * https://github.com/luben/zstd-jni +| +| This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression +| and decompression library written by William Kinney. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jfastlz.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jfastlz/ +| +| This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data +| interchange format, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.protobuf.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/protobuf +| +| This product optionally depends on 'Bouncy Castle Crypto APIs' to generate +| a temporary self-signed X.509 certificate when the JVM does not provide the +| equivalent functionality. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.bouncycastle.txt (MIT License) +| * HOMEPAGE: +| * https://www.bouncycastle.org/ +| +| This product optionally depends on 'Snappy', a compression library produced +| by Google Inc, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.snappy.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/snappy +| +| This product optionally depends on 'JBoss Marshalling', an alternative Java +| serialization API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jboss-marshalling.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jboss-remoting/jboss-marshalling +| +| This product optionally depends on 'Caliper', Google's micro- +| benchmarking framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.caliper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/google/caliper +| +| This product optionally depends on 'Apache Commons Logging', a logging +| framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-logging.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/logging/ +| +| This product optionally depends on 'Apache Log4J', a logging framework, which +| can be obtained at: +| +| * LICENSE: +| * license/LICENSE.log4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://logging.apache.org/log4j/ +| +| This product optionally depends on 'Aalto XML', an ultra-high performance +| non-blocking XML processor, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.aalto-xml.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://wiki.fasterxml.com/AaltoHome +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hpack.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/twitter/hpack +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Cory Benfield. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hyper-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/python-hyper/hpack/ +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Tatsuhiro Tsujikawa. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.nghttp2-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/nghttp2/nghttp2/ +| +| This product contains a modified portion of 'Apache Commons Lang', a Java library +| provides utilities for the java.lang API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-lang.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/proper/commons-lang/ +| +| +| This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build. +| +| * LICENSE: +| * license/LICENSE.mvn-wrapper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/takari/maven-wrapper +| +| This product contains the dnsinfo.h header file, that provides a way to retrieve the system DNS configuration on MacOS. +| This private header is also used by Apple's open source +| mDNSResponder (https://opensource.apple.com/tarballs/mDNSResponder/). +| +| * LICENSE: +| * license/LICENSE.dnsinfo.txt (Apple Public Source License 2.0) +| * HOMEPAGE: +| * https://www.opensource.apple.com/source/configd/configd-453.19/dnsinfo/dnsinfo.h +| +| This product optionally depends on 'Brotli4j', Brotli compression and +| decompression for Java., which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.brotli4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/hyperxpro/Brotli4j + +--------------------------------------------------------- + +This product includes gRPC 1.71.0, with the following in its NOTICE: + +| Copyright 2014 The gRPC Authors +| +| Licensed under the Apache License, Version 2.0 (the "License"); +| you may not use this file except in compliance with the License. +| You may obtain a copy of the License at +| +| http://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, +| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +| See the License for the specific language governing permissions and +| limitations under the License. +| +| ----------------------------------------------------------------------- +| +| This product contains a modified portion of 'OkHttp', an open source +| HTTP & SPDY client for Android and Java applications, which can be obtained +| at: +| +| * LICENSE: +| * okhttp/third_party/okhttp/LICENSE (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/square/okhttp +| * LOCATION_IN_GRPC: +| * okhttp/third_party/okhttp +| +| This product contains a modified portion of 'Envoy', an open source +| cloud-native high-performance edge/middle/service proxy, which can be +| obtained at: +| +| * LICENSE: +| * xds/third_party/envoy/LICENSE (Apache License 2.0) +| * NOTICE: +| * xds/third_party/envoy/NOTICE +| * HOMEPAGE: +| * https://www.envoyproxy.io +| * LOCATION_IN_GRPC: +| * xds/third_party/envoy +| +| This product contains a modified portion of 'protoc-gen-validate (PGV)', +| an open source protoc plugin to generate polyglot message validators, +| which can be obtained at: +| +| * LICENSE: +| * xds/third_party/protoc-gen-validate/LICENSE (Apache License 2.0) +| * NOTICE: +| * xds/third_party/protoc-gen-validate/NOTICE +| * HOMEPAGE: +| * https://github.com/envoyproxy/protoc-gen-validate +| * LOCATION_IN_GRPC: +| * xds/third_party/protoc-gen-validate +| +| This product contains a modified portion of 'udpa', +| an open source universal data plane API, which can be obtained at: +| +| * LICENSE: +| * xds/third_party/udpa/LICENSE (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/cncf/udpa +| * LOCATION_IN_GRPC: +| * xds/third_party/udpa diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml index 8ad7e76801..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 @@ -47,6 +47,21 @@ under the License. + + io.grpc + grpc-api + + + + io.grpc + grpc-netty + + + + io.netty + netty-transport + + org.apache.arrow arrow-memory-core @@ -90,14 +105,43 @@ under the License. commons-io commons-io - 2.17.0 + 2.22.0 test org.mockito mockito-core - ${mockito.core.version} + test + + + org.mockito + mockito-junit-jupiter + 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 @@ -115,19 +159,32 @@ under the License. org.apache.calcite.avatica avatica - 1.25.0 + 1.27.0 org.bouncycastle bcpkix-jdk18on - 1.79 + 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 c1b1c8f8e6..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,6 +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 { @@ -163,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(); @@ -172,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/calendar/ArrowFlightJdbcDateVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java index ebe4016209..cdafeffc32 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcDateVectorAccessor.java @@ -24,7 +24,9 @@ import static org.apache.calcite.avatica.util.DateTimeUtils.unixDateToString; import java.sql.Date; +import java.sql.SQLException; import java.sql.Timestamp; +import java.time.LocalDate; import java.util.Calendar; import java.util.concurrent.TimeUnit; import java.util.function.IntSupplier; @@ -85,6 +87,19 @@ public Object getObject() { return this.getDate(null); } + @Override + public T getObject(final Class type) throws SQLException { + final Object value; + if (type == LocalDate.class) { + value = getLocalDate(); + } else if (type == Date.class) { + value = getObject(); + } else { + throw new SQLException("Object type not supported for Date Vector"); + } + return !type.isPrimitive() && wasNull ? null : type.cast(value); + } + @Override public Date getDate(Calendar calendar) { fillHolder(); @@ -134,4 +149,8 @@ protected static TimeUnit getTimeUnitForVector(ValueVector vector) { throw new IllegalArgumentException("Invalid Arrow vector"); } + + private LocalDate getLocalDate() { + return getDate(null).toLocalDate(); + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java index debdd0fcb4..813fbc7cfd 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessor.java @@ -21,11 +21,18 @@ import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorGetter.createGetter; import java.sql.Date; +import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; +import java.time.Instant; import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; import java.time.temporal.ChronoUnit; import java.util.Calendar; +import java.util.Objects; +import java.util.Set; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import java.util.function.IntSupplier; @@ -43,6 +50,7 @@ public class ArrowFlightJdbcTimeStampVectorAccessor extends ArrowFlightJdbcAcces private final TimeUnit timeUnit; private final LongToLocalDateTime longToLocalDateTime; private final Holder holder; + private final boolean isZoned; /** Functional interface used to convert a number (in any time resolution) to LocalDateTime. */ interface LongToLocalDateTime { @@ -58,6 +66,9 @@ public ArrowFlightJdbcTimeStampVectorAccessor( this.holder = new Holder(); this.getter = createGetter(vector); + // whether the vector included TZ info + this.isZoned = getVectorIsZoned(vector); + // non-null, either the vector TZ or default to UTC this.timeZone = getTimeZoneForVector(vector); this.timeUnit = getTimeUnitForVector(vector); this.longToLocalDateTime = getLongToLocalDateTimeForVector(vector, this.timeZone); @@ -68,11 +79,62 @@ public Class getObjectClass() { return Timestamp.class; } + @Override + public T getObject(final Class type) throws SQLException { + final Object value; + if (!this.isZoned + & Set.of(OffsetDateTime.class, ZonedDateTime.class, Instant.class).contains(type)) { + throw new SQLException( + "Vectors without timezones can't be converted to objects with offset/tz info."); + } else if (type == OffsetDateTime.class) { + value = getOffsetDateTime(); + } else if (type == LocalDateTime.class) { + value = getLocalDateTime(null); + } else if (type == ZonedDateTime.class) { + value = getZonedDateTime(); + } else if (type == Instant.class) { + value = getInstant(); + } else if (type == Timestamp.class) { + value = getObject(); + } else { + throw new SQLException("Object type not supported for TimeStamp Vector"); + } + + return !type.isPrimitive() && wasNull ? null : type.cast(value); + } + @Override public Object getObject() { return this.getTimestamp(null); } + private ZonedDateTime getZonedDateTime() { + LocalDateTime localDateTime = getLocalDateTime(null); + if (localDateTime == null) { + return null; + } + + return localDateTime.atZone(this.timeZone.toZoneId()); + } + + private OffsetDateTime getOffsetDateTime() { + LocalDateTime localDateTime = getLocalDateTime(null); + if (localDateTime == null) { + return null; + } + ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime); + return localDateTime.atOffset(offset); + } + + private Instant getInstant() { + LocalDateTime localDateTime = getLocalDateTime(null); + if (localDateTime == null) { + return null; + } + ZoneOffset offset = this.timeZone.toZoneId().getRules().getOffset(localDateTime); + return localDateTime.toInstant(offset); + } + private LocalDateTime getLocalDateTime(Calendar calendar) { getter.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; @@ -85,7 +147,9 @@ private LocalDateTime getLocalDateTime(Calendar calendar) { LocalDateTime localDateTime = this.longToLocalDateTime.fromLong(value); - if (calendar != null) { + // Adjust timestamp to desired calendar (if provided) only if the column includes TZ info, + // otherwise treat as wall-clock time + if (calendar != null && this.isZoned) { TimeZone timeZone = calendar.getTimeZone(); long millis = this.timeUnit.toMillis(value); localDateTime = @@ -102,7 +166,7 @@ public Date getDate(Calendar calendar) { return null; } - return new Date(Timestamp.valueOf(localDateTime).getTime()); + return new Date(getTimestampWithOffset(calendar, localDateTime).getTime()); } @Override @@ -112,7 +176,7 @@ public Time getTime(Calendar calendar) { return null; } - return new Time(Timestamp.valueOf(localDateTime).getTime()); + return new Time(getTimestampWithOffset(calendar, localDateTime).getTime()); } @Override @@ -122,6 +186,24 @@ public Timestamp getTimestamp(Calendar calendar) { return null; } + return getTimestampWithOffset(calendar, localDateTime); + } + + /** + * Apply offset to LocalDateTime to get a Timestamp with legacy behavior. Previously we applied + * the offset to the LocalDateTime even if the underlying Vector did not have a TZ. In order to + * support java.time.* accessors, we fixed this so we only apply the offset if the underlying + * vector includes TZ info. In order to maintain backward compatibility, we apply the offset if + * needed for getDate, getTime, and getTimestamp. + */ + private Timestamp getTimestampWithOffset(Calendar calendar, LocalDateTime localDateTime) { + if (calendar != null && !isZoned) { + TimeZone timeZone = calendar.getTimeZone(); + long millis = Timestamp.valueOf(localDateTime).getTime(); + localDateTime = + localDateTime.minus( + timeZone.getOffset(millis) - this.timeZone.getOffset(millis), ChronoUnit.MILLIS); + } return Timestamp.valueOf(localDateTime); } @@ -170,11 +252,14 @@ protected static TimeZone getTimeZoneForVector(TimeStampVector vector) { ArrowType.Timestamp arrowType = (ArrowType.Timestamp) vector.getField().getFieldType().getType(); - String timezoneName = arrowType.getTimezone(); - if (timezoneName == null) { - return TimeZone.getTimeZone("UTC"); - } - + String timezoneName = Objects.requireNonNullElse(arrowType.getTimezone(), "UTC"); return TimeZone.getTimeZone(timezoneName); } + + protected static boolean getVectorIsZoned(TimeStampVector vector) { + ArrowType.Timestamp arrowType = + (ArrowType.Timestamp) vector.getField().getFieldType().getType(); + + return arrowType.getTimezone() != null; + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java index 2c03ee631e..d525c2fdd2 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeVectorAccessor.java @@ -20,8 +20,10 @@ import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.Holder; import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeVectorGetter.createGetter; +import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; +import java.time.LocalTime; import java.util.Calendar; import java.util.concurrent.TimeUnit; import java.util.function.IntSupplier; @@ -121,6 +123,19 @@ public Object getObject() { return this.getTime(null); } + @Override + public T getObject(final Class type) throws SQLException { + final Object value; + if (type == LocalTime.class) { + value = getLocalTime(); + } else if (type == Time.class) { + value = getObject(); + } else { + throw new SQLException("Object type not supported for Time Vector"); + } + return !type.isPrimitive() && wasNull ? null : type.cast(value); + } + @Override public Time getTime(Calendar calendar) { fillHolder(); @@ -134,6 +149,10 @@ public Time getTime(Calendar calendar) { return new ArrowFlightJdbcTime(DateTimeUtils.applyCalendarOffset(milliseconds, calendar)); } + private LocalTime getLocalTime() { + return getTime(null).toLocalTime(); + } + private void fillHolder() { getter.get(getCurrentRow(), holder); this.wasNull = holder.isSet == 0; 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 0e9c79a090..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 @@ -17,10 +17,13 @@ package org.apache.arrow.driver.jdbc.client; import com.google.common.collect.ImmutableMap; +import io.grpc.netty.NettyChannelBuilder; +import io.netty.channel.ChannelOption; import java.io.IOException; import java.net.URI; import java.security.GeneralSecurityException; import java.sql.SQLException; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -29,19 +32,24 @@ 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; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightClientMiddleware; import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightGrpcUtils; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightRuntimeException; 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; @@ -50,6 +58,7 @@ import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; import org.apache.arrow.flight.client.ClientCookieMiddleware; import org.apache.arrow.flight.grpc.CredentialCallOption; +import org.apache.arrow.flight.grpc.NettyClientBuilder; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo; import org.apache.arrow.flight.sql.util.TableRef; @@ -59,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; @@ -70,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; } /** @@ -96,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; } @@ -130,23 +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 location is the same as the - // original connection's - // Location and skip creating a FlightClient in that scenario. + // 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 + // 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 = @@ -158,7 +187,9 @@ public List getStreams(final FlightInfo flightInfo) new Builder(ArrowFlightSqlClientHandler.this.builder) .withHost(endpointUri.getHost()) .withPort(endpointUri.getPort()) - .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS)); + .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS)) + .withClientCache(flightClientCache) + .withConnectTimeout(builder.connectTimeout); ArrowFlightSqlClientHandler endpointHandler = null; try { @@ -172,11 +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) { @@ -221,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 { /** @@ -268,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 @@ -277,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); } } @@ -321,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; } @@ -340,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"); } } }; @@ -508,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; @@ -543,7 +693,14 @@ public static final class Builder { @VisibleForTesting Optional catalog = Optional.empty(); - // These two middleware are for internal use within build() and should not be exposed by builder + @VisibleForTesting @Nullable FlightClientCache flightClientCache; + + @VisibleForTesting @Nullable Duration connectTimeout; + + @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 @@ -553,6 +710,8 @@ public static final class Builder { @VisibleForTesting ClientCookieMiddleware.Factory cookieFactory = new ClientCookieMiddleware.Factory(); + DriverVersion driverVersion; + public Builder() {} /** @@ -579,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; @@ -587,6 +747,8 @@ public Builder() {} if (original.retainAuth) { this.authFactory = original.authFactory; } + + this.driverVersion = original.driverVersion; } /** @@ -825,6 +987,50 @@ 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) { + return Location.forGrpcTls(host, port); + } + return Location.forGrpcInsecure(host, port); + } + /** * Builds a new {@link ArrowFlightSqlClientHandler} from the provided fields. * @@ -832,7 +1038,8 @@ public Builder withCatalog(@Nullable final String catalog) { * @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); @@ -840,22 +1047,26 @@ 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); } - final FlightClient.Builder clientBuilder = FlightClient.builder().allocator(allocator); + 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); - Location location; if (useEncryption) { - location = Location.forGrpcTls(host, port); clientBuilder.useTls(); - } else { - location = Location.forGrpcInsecure(host, port); } + Location location = getLocation(); clientBuilder.location(location); if (useEncryption) { @@ -883,11 +1094,27 @@ public ArrowFlightSqlClientHandler build() throws SQLException { } } - client = clientBuilder.build(); + NettyChannelBuilder channelBuilder = clientBuilder.build(); + + channelBuilder.userAgent(userAgent); + + if (connectTimeout != null) { + channelBuilder.withOption( + ChannelOption.CONNECT_TIMEOUT_MILLIS, (int) connectTimeout.toMillis()); + } + client = + 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()); @@ -905,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/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java new file mode 100644 index 0000000000..aceadb327b --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/oauth/OAuthTokenException.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.client.oauth; + +/** + * 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); + } + + 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 e8bae2a207..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,12 +16,15 @@ */ package org.apache.arrow.driver.jdbc.utils; +import java.sql.SQLException; +import java.time.Duration; import java.util.Arrays; import java.util.HashMap; import java.util.Map; 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; @@ -30,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 { @@ -163,6 +167,21 @@ public String getCatalog() { return ArrowFlightConnectionProperty.CATALOG.getString(properties); } + /** The initial connect timeout. */ + public Duration getConnectTimeout() { + Integer timeout = ArrowFlightConnectionProperty.CONNECT_TIMEOUT_MILLIS.getInteger(properties); + if (timeout == null) { + return Duration.ofMillis( + (int) ArrowFlightConnectionProperty.CONNECT_TIMEOUT_MILLIS.defaultValue()); + } + 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}. * @@ -195,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), @@ -213,7 +264,27 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty { TOKEN("token", null, Type.STRING, false), RETAIN_COOKIES("retainCookies", true, Type.BOOLEAN, false), RETAIN_AUTH("retainAuth", true, Type.BOOLEAN, false), - CATALOG("catalog", null, Type.STRING, 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; private final Object defaultValue; 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 96cb056db2..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 @@ -16,14 +16,17 @@ */ package org.apache.arrow.driver.jdbc.utils; +import com.google.common.base.Strings; 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 { @@ -106,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; @@ -120,7 +128,12 @@ public static int getSqlTypeIdFromArrowType(ArrowType arrowType) { case Time: return Types.TIME; case Timestamp: - return Types.TIMESTAMP; + String tz = ((ArrowType.Timestamp) arrowType).getTimezone(); + if (Strings.isNullOrEmpty(tz)) { + return Types.TIMESTAMP; + } else { + return Types.TIMESTAMP_WITH_TIMEZONE; + } case Bool: return Types.BOOLEAN; case Decimal: 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 88a172e4f2..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; @@ -299,8 +301,9 @@ public class ArrowDatabaseMetadataTest { private static Connection connection; static { - List expectedGetColumnsDataTypes = Arrays.asList(3, 93, 4); - List expectedGetColumnsTypeName = Arrays.asList("DECIMAL", "TIMESTAMP", "INTEGER"); + List expectedGetColumnsDataTypes = Arrays.asList(3, 2014, 4); + List expectedGetColumnsTypeName = + Arrays.asList("DECIMAL", "TIMESTAMP_WITH_TIMEZONE", "INTEGER"); List expectedGetColumnsRadix = Arrays.asList(10, null, 10); List expectedGetColumnsColumnSize = Arrays.asList(5, 29, 10); List expectedGetColumnsDecimalDigits = Arrays.asList(2, 9, 0); @@ -321,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, @@ -418,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"); @@ -1513,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 a8d04dfc83..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,21 +22,20 @@ 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.assertArrayEquals; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -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.fail; +import static org.junit.jupiter.api.Assertions.*; import com.google.common.collect.ImmutableSet; 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; @@ -47,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; @@ -66,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; @@ -645,6 +646,138 @@ public void testFallbackSecondFlightServer() throws Exception { } } + @Test + public void testFallbackUnresolvableFlightServer() throws Exception { + final Schema schema = + new Schema( + Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType()))); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) { + resultData.setRowCount(1); + ((IntVector) resultData.getVector(0)).set(0, 1); + + try (final FallbackFlightSqlProducer rootProducer = + new FallbackFlightSqlProducer(resultData); + FlightServer rootServer = + FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer) + .build() + .start(); + Connection newConnection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false", + rootServer.getLocation().getUri().getHost(), rootServer.getPort()))) { + // This first attempt should take a measurable amount of time. + long start = System.nanoTime(); + try (Statement newStatement = newConnection.createStatement()) { + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { + List actualData = new ArrayList<>(); + while (result.next()) { + actualData.add(result.getInt(1)); + } + + // Assert + assertEquals(resultData.getRowCount(), actualData.size()); + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); + } + } + long attempt1 = System.nanoTime(); + double elapsedMs = (attempt1 - start) / 1_000_000.; + assertTrue( + elapsedMs >= 5000., + String.format( + "Expected first attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); + + // Once the client cache is implemented (GH-661), this second attempt should take less time, + // since the failure from before should be cached. + start = System.nanoTime(); + try (Statement newStatement = newConnection.createStatement()) { + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { + List actualData = new ArrayList<>(); + while (result.next()) { + actualData.add(result.getInt(1)); + } + + // Assert + assertEquals(resultData.getRowCount(), actualData.size()); + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); + } + } + attempt1 = System.nanoTime(); + elapsedMs = (attempt1 - start) / 1_000_000.; + assertTrue( + elapsedMs < 5000., + String.format("Expected second attempt to be faster, but %f ms elapsed", elapsedMs)); + } + } + } + + @Test + public void testFallbackUnresolvableFlightServerDisableCache() throws Exception { + final Schema schema = + new Schema( + Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType()))); + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) { + resultData.setRowCount(1); + ((IntVector) resultData.getVector(0)).set(0, 1); + + try (final FallbackFlightSqlProducer rootProducer = + new FallbackFlightSqlProducer(resultData); + FlightServer rootServer = + FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer) + .build() + .start(); + Connection newConnection = + DriverManager.getConnection( + String.format( + "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false&useClientCache=false", + rootServer.getLocation().getUri().getHost(), rootServer.getPort()))) { + // This first attempt should take a measurable amount of time. + long start = System.nanoTime(); + try (Statement newStatement = newConnection.createStatement()) { + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { + List actualData = new ArrayList<>(); + while (result.next()) { + actualData.add(result.getInt(1)); + } + + // Assert + assertEquals(resultData.getRowCount(), actualData.size()); + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); + } + } + long attempt1 = System.nanoTime(); + double elapsedMs = (attempt1 - start) / 1_000_000.; + assertTrue( + elapsedMs >= 5000., + String.format( + "Expected first attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); + + // This second attempt should take a long time still, since we disabled the cache. + start = System.nanoTime(); + try (Statement newStatement = newConnection.createStatement()) { + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { + List actualData = new ArrayList<>(); + while (result.next()) { + actualData.add(result.getInt(1)); + } + + // Assert + assertEquals(resultData.getRowCount(), actualData.size()); + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); + } + } + attempt1 = System.nanoTime(); + elapsedMs = (attempt1 - start) / 1_000_000.; + assertTrue( + elapsedMs >= 5000., + String.format( + "Expected second attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); + } + } + } + @Test public void testShouldRunSelectQueryWithEmptyVectorsEmbedded() throws Exception { try (Statement statement = connection.createStatement(); @@ -668,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/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java index 2e329f148e..e4863bd80e 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/accessor/impl/calendar/ArrowFlightJdbcTimeStampVectorAccessorTest.java @@ -16,17 +16,23 @@ */ package org.apache.arrow.driver.jdbc.accessor.impl.calendar; -import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.getTimeUnitForVector; -import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.getTimeZoneForVector; +import static org.apache.arrow.driver.jdbc.accessor.impl.calendar.ArrowFlightJdbcTimeStampVectorAccessor.*; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.Date; +import java.sql.SQLException; import java.sql.Time; import java.sql.Timestamp; +import java.time.Instant; import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneId; +import java.time.ZonedDateTime; import java.util.Calendar; +import java.util.Objects; import java.util.TimeZone; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; @@ -199,6 +205,99 @@ public void testShouldGetTimestampReturnValidTimestampWithCalendar( }); } + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidLocalDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + final LocalDateTime value = accessor.getObject(LocalDateTime.class); + final LocalDateTime expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toLocalDateTime(); + + assertThat(value, equalTo(expectedValue)); + assertThat(accessor.wasNull(), is(false)); + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidInstant( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final Instant value = accessor.getObject(Instant.class); + final Instant expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toInstant(); + + assertThat(value, equalTo(expectedValue)); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(Instant.class)); + } + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidOffsetDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final OffsetDateTime value = accessor.getObject(OffsetDateTime.class); + final OffsetDateTime expectedValue = + getZonedDateTime(currentRow, expectedTimeZone).toOffsetDateTime(); + + assertThat(value, equalTo(expectedValue)); + assertThat(value.getOffset(), equalTo(expectedValue.getOffset())); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(OffsetDateTime.class)); + } + }); + } + + @ParameterizedTest + @MethodSource("data") + public void testShouldGetObjectReturnValidZonedDateTime( + Supplier vectorSupplier, String vectorType, String timeZone) + throws Exception { + setup(vectorSupplier); + final String expectedTimeZone = Objects.requireNonNullElse(timeZone, "UTC"); + final boolean vectorHasTz = timeZone != null; + accessorIterator.iterate( + vector, + (accessor, currentRow) -> { + if (vectorHasTz) { + final ZonedDateTime value = accessor.getObject(ZonedDateTime.class); + final ZonedDateTime expectedValue = getZonedDateTime(currentRow, expectedTimeZone); + + assertThat(value, equalTo(expectedValue)); + assertThat(value.getZone(), equalTo(ZoneId.of(expectedTimeZone))); + assertThat(accessor.wasNull(), is(false)); + } else { + assertThrows(SQLException.class, () -> accessor.getObject(ZonedDateTime.class)); + } + }); + } + @ParameterizedTest @MethodSource("data") public void testShouldGetTimestampReturnNull(Supplier vectorSupplier) { @@ -320,6 +419,21 @@ private Timestamp getTimestampForVector(int currentRow, String timeZone) { return expectedTimestamp; } + /** ZonedDateTime contains all necessary information to generate any java.time object. */ + private ZonedDateTime getZonedDateTime(int currentRow, String timeZone) { + Object object = vector.getObject(currentRow); + TimeZone tz = TimeZone.getTimeZone(timeZone); + ZonedDateTime expectedTimestamp = null; + if (object instanceof LocalDateTime) { + expectedTimestamp = ((LocalDateTime) object).atZone(tz.toZoneId()); + } else if (object instanceof Long) { + TimeUnit timeUnit = getTimeUnitForVector(vector); + Instant instant = Instant.ofEpochMilli(timeUnit.toMillis((Long) object)); + expectedTimestamp = ZonedDateTime.ofInstant(instant, tz.toZoneId()); + } + return expectedTimestamp; + } + @ParameterizedTest @MethodSource("data") public void testShouldGetObjectClass(Supplier vectorSupplier) throws Exception { 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 6beaba8236..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,6 +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 4a46b5f5be..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 @@ -18,15 +18,18 @@ import static java.lang.Runtime.getRuntime; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.CATALOG; +import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.CONNECT_TIMEOUT_MILLIS; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.HOST; import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.PASSWORD; 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; +import java.time.Duration; import java.util.Properties; import java.util.Random; import java.util.function.Function; @@ -59,49 +62,73 @@ public void setUp() { public void testGetProperty( ArrowFlightConnectionProperty property, Object value, + Object expected, Function configFunction) { properties.put(property.camelName(), value); arrowFlightConnectionConfigFunction = configFunction; - assertThat(configFunction.apply(arrowFlightConnectionConfig), is(value)); - assertThat(arrowFlightConnectionConfigFunction.apply(arrowFlightConnectionConfig), is(value)); + assertThat(configFunction.apply(arrowFlightConnectionConfig), is(expected)); + assertThat( + arrowFlightConnectionConfigFunction.apply(arrowFlightConnectionConfig), is(expected)); } public static Stream provideParameters() { + int port = RANDOM.nextInt(Short.toUnsignedInt(Short.MAX_VALUE)); + boolean useEncryption = RANDOM.nextBoolean(); + int threadPoolSize = RANDOM.nextInt(getRuntime().availableProcessors()); return Stream.of( Arguments.of( HOST, "host", + "host", (Function) ArrowFlightConnectionConfigImpl::getHost), Arguments.of( PORT, - RANDOM.nextInt(Short.toUnsignedInt(Short.MAX_VALUE)), + port, + port, (Function) ArrowFlightConnectionConfigImpl::getPort), Arguments.of( USER, "user", + "user", (Function) ArrowFlightConnectionConfigImpl::getUser), Arguments.of( PASSWORD, "password", + "password", (Function) ArrowFlightConnectionConfigImpl::getPassword), Arguments.of( USE_ENCRYPTION, - RANDOM.nextBoolean(), + useEncryption, + useEncryption, (Function) ArrowFlightConnectionConfigImpl::useEncryption), Arguments.of( THREAD_POOL_SIZE, - RANDOM.nextInt(getRuntime().availableProcessors()), + threadPoolSize, + threadPoolSize, (Function) ArrowFlightConnectionConfigImpl::threadPoolSize), Arguments.of( CATALOG, "catalog", + "catalog", + (Function) + ArrowFlightConnectionConfigImpl::getCatalog), + Arguments.of( + CONNECT_TIMEOUT_MILLIS, + 5000, + Duration.ofMillis(5000), + (Function) + ArrowFlightConnectionConfigImpl::getConnectTimeout), + Arguments.of( + USE_CLIENT_CACHE, + false, + false, (Function) - ArrowFlightConnectionConfigImpl::getCatalog)); + 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/FallbackFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FallbackFlightSqlProducer.java index 9aa257172c..670b9e3be0 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FallbackFlightSqlProducer.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FallbackFlightSqlProducer.java @@ -109,6 +109,16 @@ private FlightInfo getFlightInfo(FlightDescriptor descriptor, String query) { Location.forGrpcInsecure("localhost", 9999), Location.reuseConnection()) .build()); + } else if (query.equals("fallback with unresolvable")) { + endpoints = + Collections.singletonList( + FlightEndpoint.builder( + ticket, + // Inaccessible IP + // https://stackoverflow.com/questions/10456044/what-is-a-good-invalid-ip-address-to-use-for-unit-tests + Location.forGrpcInsecure("203.0.113.0", 9999), + Location.reuseConnection()) + .build()); } else { throw CallStatus.UNIMPLEMENTED.withDescription(query).toRuntimeException(); } 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 00af3c96ba..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,17 +41,25 @@ 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))); assertEquals( Types.TIME, getSqlTypeIdFromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + assertEquals( + Types.TIMESTAMP, + getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( Types.TIMESTAMP, getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, ""))); + assertEquals( + Types.TIMESTAMP_WITH_TIMEZONE, + getSqlTypeIdFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))); assertEquals(Types.BOOLEAN, getSqlTypeIdFromArrowType(new ArrowType.Bool())); @@ -77,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 @@ -88,16 +99,24 @@ 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))); assertEquals("TIME", getSqlTypeNameFromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + assertEquals( + "TIMESTAMP", + getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, null))); assertEquals( "TIMESTAMP", getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, ""))); + assertEquals( + "TIMESTAMP_WITH_TIMEZONE", + getSqlTypeNameFromArrowType(new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"))); assertEquals("BOOLEAN", getSqlTypeNameFromArrowType(new ArrowType.Bool())); @@ -124,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 ae8c543fbf..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,15 @@ under the License. - + + + META-INF/LICENSE.txt + src/shade/LICENSE.txt + + + META-INF/NOTICE.txt + src/shade/NOTICE.txt + @@ -151,6 +159,7 @@ under the License. org.apache.calcite.avatica:* META-INF/services/java.sql.Driver + META-INF/README.txt @@ -166,6 +175,13 @@ under the License. META-INF/versions/ **/*.proto **/module-info.class + + LICENSE.txt + NOTICE.txt + META-INF/*LICENSE* + META-INF/*NOTICE* + META-INF/license/* + META-INF/licenses/**/* diff --git a/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt new file mode 100644 index 0000000000..8476bd9995 --- /dev/null +++ b/flight/flight-sql-jdbc-driver/src/shade/LICENSE.txt @@ -0,0 +1,375 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-------------------------------------------------------------------------------- + +This binary artifact contains gRPC 1.71.0. + +Copyright: Copyright 2014 The gRPC Authors +Home page: https://grpc.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Netty 4.1.119.Final. + +Copyright: Copyright 2014 The Netty Project +Home page: https://netty.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Gson 2.11.0. + +Copyright: Copyright 2008 Google Inc. +Home page: https://github.com/google/gson +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +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.2. + +Copyright: Copyright 2008 Google Inc. All rights reserved. +Home page: https://protobuf.dev/ +License: https://github.com/protocolbuffers/protobuf/blob/v4.30.1/LICENSE (BSD) +License text: + +| Copyright 2008 Google Inc. All rights reserved. +| +| Redistribution and use in source and binary forms, with or without +| modification, are permitted provided that the following conditions are +| met: +| +| * Redistributions of source code must retain the above copyright +| notice, this list of conditions and the following disclaimer. +| * Redistributions in binary form must reproduce the above +| copyright notice, this list of conditions and the following disclaimer +| in the documentation and/or other materials provided with the +| distribution. +| * Neither the name of Google Inc. nor the names of its +| contributors may be used to endorse or promote products derived from +| this software without specific prior written permission. +| +| THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +| "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +| LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +| A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +| OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +| SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +| LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +| DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +| THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +| (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +| OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +| +| Code generated by the Protocol Buffer compiler is owned by the owner +| of the input file used when generating it. This code is not +| standalone and requires a support library to be linked with it. This +| support library is itself covered by the above license. + +-------------------------------------------------------------------------------- + +This binary artifact contains Jackson 2.18.3. + +Home page: https://github.com/FasterXML/jackson +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons Codec 1.18.0. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-codec/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Google Flatbuffers 25.2.10. + +Home page: https://flatbuffers.dev/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Slf4j 2.0.17. + +Copyright: Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +Home page: http://www.slf4j.org/ +License: MIT +License text: + +| Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +| All rights reserved. +| +| Permission is hereby granted, free of charge, to any person obtaining +| a copy of this software and associated documentation files (the +| "Software"), to deal in the Software without restriction, including +| without limitation the rights to use, copy, modify, merge, publish, +| distribute, sublicense, and/or sell copies of the Software, and to +| permit persons to whom the Software is furnished to do so, subject to +| the following conditions: +| +| The above copyright notice and this permission notice shall be +| included in all copies or substantial portions of the Software. +| +| THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +| EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +| MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +| NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +| LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +| OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +| WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Calcite Avatica 1.26.0. + +Copyright: Copyright 2012-2024 The Apache Software Foundation +Home page: https://calcite.apache.org/avatica/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache HttpComponents HttpClient 4.5.13. + +Copyright: Copyright 1999-2020 The Apache Software Foundation +Home page: https://hc.apache.org/index.html +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). +Home page: https://www.bouncycastle.org/ +License: MIT +License text: + +| Copyright (c) 2000-2024 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org). +| Permission is hereby granted, free of charge, to any person obtaining a copy of this software and +| associated documentation files (the "Software"), to deal in the Software without restriction, +| including without limitation the rights to use, copy, modify, merge, publish, distribute, +| sub license, and/or sell copies of the Software, and to permit persons to whom the Software is +| furnished to do so, subject to the following conditions: The above copyright notice and this +| permission notice shall be included in all copies or substantial portions of the Software. +| +| **THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT +| NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +| NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +| DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT +| OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.** diff --git a/flight/flight-sql-jdbc-driver/src/shade/NOTICE.txt b/flight/flight-sql-jdbc-driver/src/shade/NOTICE.txt new file mode 100644 index 0000000000..2f2f6b9e75 --- /dev/null +++ b/flight/flight-sql-jdbc-driver/src/shade/NOTICE.txt @@ -0,0 +1,340 @@ +Apache Arrow Java +Copyright 2016-2025 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +--------------------------------------------------- + +This product includes gRPC 1.71.0, with the following in its NOTICE: + +| Copyright 2014 The gRPC Authors +| +| Licensed under the Apache License, Version 2.0 (the "License"); +| you may not use this file except in compliance with the License. +| You may obtain a copy of the License at +| +| http://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, +| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +| See the License for the specific language governing permissions and +| limitations under the License. +| +| ----------------------------------------------------------------------- +| +| This product contains a modified portion of 'OkHttp', an open source +| HTTP & SPDY client for Android and Java applications, which can be obtained +| at: +| +| * LICENSE: +| * okhttp/third_party/okhttp/LICENSE (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/square/okhttp +| * LOCATION_IN_GRPC: +| * okhttp/third_party/okhttp +| +| This product contains a modified portion of 'Envoy', an open source +| cloud-native high-performance edge/middle/service proxy, which can be +| obtained at: +| +| * LICENSE: +| * xds/third_party/envoy/LICENSE (Apache License 2.0) +| * NOTICE: +| * xds/third_party/envoy/NOTICE +| * HOMEPAGE: +| * https://www.envoyproxy.io +| * LOCATION_IN_GRPC: +| * xds/third_party/envoy +| +| This product contains a modified portion of 'protoc-gen-validate (PGV)', +| an open source protoc plugin to generate polyglot message validators, +| which can be obtained at: +| +| * LICENSE: +| * xds/third_party/protoc-gen-validate/LICENSE (Apache License 2.0) +| * NOTICE: +| * xds/third_party/protoc-gen-validate/NOTICE +| * HOMEPAGE: +| * https://github.com/envoyproxy/protoc-gen-validate +| * LOCATION_IN_GRPC: +| * xds/third_party/protoc-gen-validate +| +| This product contains a modified portion of 'udpa', +| an open source universal data plane API, which can be obtained at: +| +| * LICENSE: +| * xds/third_party/udpa/LICENSE (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/cncf/udpa +| * LOCATION_IN_GRPC: +| * xds/third_party/udpa + +--------------------------------------------------- + +This product includes Netty 4.1.119.Final, with the following in its NOTICE: + +| The Netty Project +| ================= +| +| Please visit the Netty web site for more information: +| +| * https://netty.io/ +| +| Copyright 2014 The Netty Project +| +| The Netty Project licenses this file to you under the Apache License, +| version 2.0 (the "License"); you may not use this file except in compliance +| with the License. You may obtain a copy of the License at: +| +| https://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +| WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +| License for the specific language governing permissions and limitations +| under the License. +| +| Also, please refer to each LICENSE..txt file, which is located in +| the 'license' directory of the distribution file, for the license terms of the +| components that this product depends on. +| +| ------------------------------------------------------------------------------- +| This product contains the extensions to Java Collections Framework which has +| been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene: +| +| * LICENSE: +| * license/LICENSE.jsr166y.txt (Public Domain) +| * HOMEPAGE: +| * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/ +| * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/ +| +| This product contains a modified version of Robert Harder's Public Domain +| Base64 Encoder and Decoder, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.base64.txt (Public Domain) +| * HOMEPAGE: +| * http://iharder.sourceforge.net/current/java/base64/ +| +| This product contains a modified portion of 'Webbit', an event based +| WebSocket and HTTP server, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.webbit.txt (BSD License) +| * HOMEPAGE: +| * https://github.com/joewalnes/webbit +| +| This product contains a modified portion of 'SLF4J', a simple logging +| facade for Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.slf4j.txt (MIT License) +| * HOMEPAGE: +| * https://www.slf4j.org/ +| +| This product contains a modified portion of 'Apache Harmony', an open source +| Java SE, which can be obtained at: +| +| * NOTICE: +| * license/NOTICE.harmony.txt +| * LICENSE: +| * license/LICENSE.harmony.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://archive.apache.org/dist/harmony/ +| +| This product contains a modified portion of 'jbzip2', a Java bzip2 compression +| and decompression library written by Matthew J. Francis. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jbzip2.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jbzip2/ +| +| This product contains a modified portion of 'libdivsufsort', a C API library to construct +| the suffix array and the Burrows-Wheeler transformed string for any input string of +| a constant-size alphabet written by Yuta Mori. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.libdivsufsort.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/y-256/libdivsufsort +| +| This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM, +| which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jctools.txt (ASL2 License) +| * HOMEPAGE: +| * https://github.com/JCTools/JCTools +| +| This product optionally depends on 'JZlib', a re-implementation of zlib in +| pure Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jzlib.txt (BSD style License) +| * HOMEPAGE: +| * http://www.jcraft.com/jzlib/ +| +| This product optionally depends on 'Compress-LZF', a Java library for encoding and +| decoding data in LZF format, written by Tatu Saloranta. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.compress-lzf.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/ning/compress +| +| This product optionally depends on 'lz4', a LZ4 Java compression +| and decompression library written by Adrien Grand. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lz4.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jpountz/lz4-java +| +| This product optionally depends on 'lzma-java', a LZMA Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lzma-java.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jponge/lzma-java +| +| This product optionally depends on 'zstd-jni', a zstd-jni Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.zstd-jni.txt (BSD) +| * HOMEPAGE: +| * https://github.com/luben/zstd-jni +| +| This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression +| and decompression library written by William Kinney. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jfastlz.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jfastlz/ +| +| This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data +| interchange format, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.protobuf.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/protobuf +| +| This product optionally depends on 'Bouncy Castle Crypto APIs' to generate +| a temporary self-signed X.509 certificate when the JVM does not provide the +| equivalent functionality. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.bouncycastle.txt (MIT License) +| * HOMEPAGE: +| * https://www.bouncycastle.org/ +| +| This product optionally depends on 'Snappy', a compression library produced +| by Google Inc, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.snappy.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/snappy +| +| This product optionally depends on 'JBoss Marshalling', an alternative Java +| serialization API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jboss-marshalling.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jboss-remoting/jboss-marshalling +| +| This product optionally depends on 'Caliper', Google's micro- +| benchmarking framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.caliper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/google/caliper +| +| This product optionally depends on 'Apache Commons Logging', a logging +| framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-logging.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/logging/ +| +| This product optionally depends on 'Apache Log4J', a logging framework, which +| can be obtained at: +| +| * LICENSE: +| * license/LICENSE.log4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://logging.apache.org/log4j/ +| +| This product optionally depends on 'Aalto XML', an ultra-high performance +| non-blocking XML processor, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.aalto-xml.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://wiki.fasterxml.com/AaltoHome +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hpack.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/twitter/hpack +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Cory Benfield. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hyper-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/python-hyper/hpack/ +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Tatsuhiro Tsujikawa. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.nghttp2-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/nghttp2/nghttp2/ +| +| This product contains a modified portion of 'Apache Commons Lang', a Java library +| provides utilities for the java.lang API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-lang.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/proper/commons-lang/ +| +| +| This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build. +| +| * LICENSE: +| * license/LICENSE.mvn-wrapper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/takari/maven-wrapper +| +| This product contains the dnsinfo.h header file, that provides a way to retrieve the system DNS configuration on MacOS. +| This private header is also used by Apple's open source +| mDNSResponder (https://opensource.apple.com/tarballs/mDNSResponder/). +| +| * LICENSE: +| * license/LICENSE.dnsinfo.txt (Apple Public Source License 2.0) +| * HOMEPAGE: +| * https://www.opensource.apple.com/source/configd/configd-453.19/dnsinfo/dnsinfo.h +| +| This product optionally depends on 'Brotli4j', Brotli compression and +| decompression for Java., which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.brotli4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/hyperxpro/Brotli4j 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 a0e108d6a0..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 @@ -67,7 +67,13 @@ public class ITDriverJarValidation { /** List of allowed files a jar entry may match. */ public static final Set ALLOWED_FILES = ImmutableSet.of( + "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 9cbc8430fe..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.12.0 + 2.14.0 test @@ -107,25 +107,25 @@ under the License. org.apache.commons commons-pool2 - 2.12.0 + 2.13.1 test org.apache.commons commons-text - 1.12.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 9a6ffdfdca..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 @@ -92,7 +92,7 @@ import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; @@ -236,7 +236,7 @@ public long executeIngest( * @return the number of rows affected. */ public long executeIngest( - final ArrowStreamReader dataReader, + final ArrowReader dataReader, final ExecuteIngestOptions ingestOptions, final CallOption... options) { return executeIngest(dataReader, ingestOptions, /*transaction*/ null, options); @@ -270,7 +270,7 @@ public long executeIngest( * @return the number of rows affected. */ public long executeIngest( - final ArrowStreamReader dataReader, + final ArrowReader dataReader, final ExecuteIngestOptions ingestOptions, Transaction transaction, final CallOption... options) { @@ -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/format/src/main/java/org/apache/arrow/flatbuf/Binary.java b/format/src/main/java/org/apache/arrow/flatbuf/Binary.java index 65b3f4f577..62938e9cd0 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Binary.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Binary.java @@ -38,7 +38,7 @@ */ @SuppressWarnings("unused") public final class Binary extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Binary getRootAsBinary(ByteBuffer _bb) { return getRootAsBinary(_bb, new Binary()); } public static Binary getRootAsBinary(ByteBuffer _bb, Binary obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/BinaryView.java b/format/src/main/java/org/apache/arrow/flatbuf/BinaryView.java index f20245ca28..4efcec0eb4 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/BinaryView.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/BinaryView.java @@ -44,7 +44,7 @@ */ @SuppressWarnings("unused") public final class BinaryView extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static BinaryView getRootAsBinaryView(ByteBuffer _bb) { return getRootAsBinaryView(_bb, new BinaryView()); } public static BinaryView getRootAsBinaryView(ByteBuffer _bb, BinaryView obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/BodyCompression.java b/format/src/main/java/org/apache/arrow/flatbuf/BodyCompression.java index a7e4288467..4b6fcd2c0b 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/BodyCompression.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/BodyCompression.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class BodyCompression extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static BodyCompression getRootAsBodyCompression(ByteBuffer _bb) { return getRootAsBodyCompression(_bb, new BodyCompression()); } public static BodyCompression getRootAsBodyCompression(ByteBuffer _bb, BodyCompression obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Bool.java b/format/src/main/java/org/apache/arrow/flatbuf/Bool.java index 7d03705ec8..0c9c4224f4 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Bool.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Bool.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Bool extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Bool getRootAsBool(ByteBuffer _bb) { return getRootAsBool(_bb, new Bool()); } public static Bool getRootAsBool(ByteBuffer _bb, Bool obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Date.java b/format/src/main/java/org/apache/arrow/flatbuf/Date.java index b9af04b120..42a518de06 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Date.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Date.java @@ -43,7 +43,7 @@ */ @SuppressWarnings("unused") public final class Date extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Date getRootAsDate(ByteBuffer _bb) { return getRootAsDate(_bb, new Date()); } public static Date getRootAsDate(ByteBuffer _bb, Date obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Decimal.java b/format/src/main/java/org/apache/arrow/flatbuf/Decimal.java index e2f38558d4..5eaadf234c 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Decimal.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Decimal.java @@ -41,7 +41,7 @@ */ @SuppressWarnings("unused") public final class Decimal extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Decimal getRootAsDecimal(ByteBuffer _bb) { return getRootAsDecimal(_bb, new Decimal()); } public static Decimal getRootAsDecimal(ByteBuffer _bb, Decimal obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/DictionaryBatch.java b/format/src/main/java/org/apache/arrow/flatbuf/DictionaryBatch.java index 891df5d164..2bee8853c7 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/DictionaryBatch.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/DictionaryBatch.java @@ -43,7 +43,7 @@ */ @SuppressWarnings("unused") public final class DictionaryBatch extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static DictionaryBatch getRootAsDictionaryBatch(ByteBuffer _bb) { return getRootAsDictionaryBatch(_bb, new DictionaryBatch()); } public static DictionaryBatch getRootAsDictionaryBatch(ByteBuffer _bb, DictionaryBatch obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/DictionaryEncoding.java b/format/src/main/java/org/apache/arrow/flatbuf/DictionaryEncoding.java index 28c52d579e..ad1f74f833 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/DictionaryEncoding.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/DictionaryEncoding.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class DictionaryEncoding extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static DictionaryEncoding getRootAsDictionaryEncoding(ByteBuffer _bb) { return getRootAsDictionaryEncoding(_bb, new DictionaryEncoding()); } public static DictionaryEncoding getRootAsDictionaryEncoding(ByteBuffer _bb, DictionaryEncoding obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Duration.java b/format/src/main/java/org/apache/arrow/flatbuf/Duration.java index 442310dbfd..bef934cef3 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Duration.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Duration.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Duration extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Duration getRootAsDuration(ByteBuffer _bb) { return getRootAsDuration(_bb, new Duration()); } public static Duration getRootAsDuration(ByteBuffer _bb, Duration obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Field.java b/format/src/main/java/org/apache/arrow/flatbuf/Field.java index 9cd50d0b98..de7b25b8db 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Field.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Field.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class Field extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Field getRootAsField(ByteBuffer _bb) { return getRootAsField(_bb, new Field()); } public static Field getRootAsField(ByteBuffer _bb, Field obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeBinary.java b/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeBinary.java index 7c4c08567f..2e3e06d580 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeBinary.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeBinary.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class FixedSizeBinary extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static FixedSizeBinary getRootAsFixedSizeBinary(ByteBuffer _bb) { return getRootAsFixedSizeBinary(_bb, new FixedSizeBinary()); } public static FixedSizeBinary getRootAsFixedSizeBinary(ByteBuffer _bb, FixedSizeBinary obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeList.java b/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeList.java index 44eceafee6..c6e2e56f22 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeList.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/FixedSizeList.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class FixedSizeList extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static FixedSizeList getRootAsFixedSizeList(ByteBuffer _bb) { return getRootAsFixedSizeList(_bb, new FixedSizeList()); } public static FixedSizeList getRootAsFixedSizeList(ByteBuffer _bb, FixedSizeList obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/FloatingPoint.java b/format/src/main/java/org/apache/arrow/flatbuf/FloatingPoint.java index 59741cde59..8f552b8f78 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/FloatingPoint.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/FloatingPoint.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class FloatingPoint extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static FloatingPoint getRootAsFloatingPoint(ByteBuffer _bb) { return getRootAsFloatingPoint(_bb, new FloatingPoint()); } public static FloatingPoint getRootAsFloatingPoint(ByteBuffer _bb, FloatingPoint obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Footer.java b/format/src/main/java/org/apache/arrow/flatbuf/Footer.java index 156215b202..61b6accd2a 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Footer.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Footer.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class Footer extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Footer getRootAsFooter(ByteBuffer _bb) { return getRootAsFooter(_bb, new Footer()); } public static Footer getRootAsFooter(ByteBuffer _bb, Footer obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Int.java b/format/src/main/java/org/apache/arrow/flatbuf/Int.java index 9dd1d448a4..378c0125eb 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Int.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Int.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Int extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Int getRootAsInt(ByteBuffer _bb) { return getRootAsInt(_bb, new Int()); } public static Int getRootAsInt(ByteBuffer _bb, Int obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Interval.java b/format/src/main/java/org/apache/arrow/flatbuf/Interval.java index 7b2be74c5c..a9fe92f3e0 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Interval.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Interval.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Interval extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Interval getRootAsInterval(ByteBuffer _bb) { return getRootAsInterval(_bb, new Interval()); } public static Interval getRootAsInterval(ByteBuffer _bb, Interval obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/KeyValue.java b/format/src/main/java/org/apache/arrow/flatbuf/KeyValue.java index 8f242ef3a9..fb35b6640a 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/KeyValue.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/KeyValue.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class KeyValue extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static KeyValue getRootAsKeyValue(ByteBuffer _bb) { return getRootAsKeyValue(_bb, new KeyValue()); } public static KeyValue getRootAsKeyValue(ByteBuffer _bb, KeyValue obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/LargeBinary.java b/format/src/main/java/org/apache/arrow/flatbuf/LargeBinary.java index 72135992a0..e74276a37c 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/LargeBinary.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/LargeBinary.java @@ -39,7 +39,7 @@ */ @SuppressWarnings("unused") public final class LargeBinary extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static LargeBinary getRootAsLargeBinary(ByteBuffer _bb) { return getRootAsLargeBinary(_bb, new LargeBinary()); } public static LargeBinary getRootAsLargeBinary(ByteBuffer _bb, LargeBinary obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/LargeList.java b/format/src/main/java/org/apache/arrow/flatbuf/LargeList.java index 1cce67712d..fe91784c82 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/LargeList.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/LargeList.java @@ -39,7 +39,7 @@ */ @SuppressWarnings("unused") public final class LargeList extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static LargeList getRootAsLargeList(ByteBuffer _bb) { return getRootAsLargeList(_bb, new LargeList()); } public static LargeList getRootAsLargeList(ByteBuffer _bb, LargeList obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/LargeListView.java b/format/src/main/java/org/apache/arrow/flatbuf/LargeListView.java index 234ce35f68..b062dc67b3 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/LargeListView.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/LargeListView.java @@ -39,7 +39,7 @@ */ @SuppressWarnings("unused") public final class LargeListView extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static LargeListView getRootAsLargeListView(ByteBuffer _bb) { return getRootAsLargeListView(_bb, new LargeListView()); } public static LargeListView getRootAsLargeListView(ByteBuffer _bb, LargeListView obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/LargeUtf8.java b/format/src/main/java/org/apache/arrow/flatbuf/LargeUtf8.java index 377bfbb2cf..f2b0f8bb3e 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/LargeUtf8.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/LargeUtf8.java @@ -39,7 +39,7 @@ */ @SuppressWarnings("unused") public final class LargeUtf8 extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static LargeUtf8 getRootAsLargeUtf8(ByteBuffer _bb) { return getRootAsLargeUtf8(_bb, new LargeUtf8()); } public static LargeUtf8 getRootAsLargeUtf8(ByteBuffer _bb, LargeUtf8 obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/List.java b/format/src/main/java/org/apache/arrow/flatbuf/List.java index 7f06dda072..eede1d4c1b 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/List.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/List.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class List extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static List getRootAsList(ByteBuffer _bb) { return getRootAsList(_bb, new List()); } public static List getRootAsList(ByteBuffer _bb, List obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/ListView.java b/format/src/main/java/org/apache/arrow/flatbuf/ListView.java index 5d87df5e4c..19415a17d9 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/ListView.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/ListView.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class ListView extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static ListView getRootAsListView(ByteBuffer _bb) { return getRootAsListView(_bb, new ListView()); } public static ListView getRootAsListView(ByteBuffer _bb, ListView obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Map.java b/format/src/main/java/org/apache/arrow/flatbuf/Map.java index 10652c0544..8f996bea45 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Map.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Map.java @@ -62,7 +62,7 @@ */ @SuppressWarnings("unused") public final class Map extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Map getRootAsMap(ByteBuffer _bb) { return getRootAsMap(_bb, new Map()); } public static Map getRootAsMap(ByteBuffer _bb, Map obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Message.java b/format/src/main/java/org/apache/arrow/flatbuf/Message.java index f2518f9046..45066eb965 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Message.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Message.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Message extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Message getRootAsMessage(ByteBuffer _bb) { return getRootAsMessage(_bb, new Message()); } public static Message getRootAsMessage(ByteBuffer _bb, Message obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Null.java b/format/src/main/java/org/apache/arrow/flatbuf/Null.java index 60196193ea..88a0c31f5a 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Null.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Null.java @@ -38,7 +38,7 @@ */ @SuppressWarnings("unused") public final class Null extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Null getRootAsNull(ByteBuffer _bb) { return getRootAsNull(_bb, new Null()); } public static Null getRootAsNull(ByteBuffer _bb, Null obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/RecordBatch.java b/format/src/main/java/org/apache/arrow/flatbuf/RecordBatch.java index 83deda7480..db68dc2150 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/RecordBatch.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/RecordBatch.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class RecordBatch extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static RecordBatch getRootAsRecordBatch(ByteBuffer _bb) { return getRootAsRecordBatch(_bb, new RecordBatch()); } public static RecordBatch getRootAsRecordBatch(ByteBuffer _bb, RecordBatch obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/RunEndEncoded.java b/format/src/main/java/org/apache/arrow/flatbuf/RunEndEncoded.java index 070ad9499e..a1b7699597 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/RunEndEncoded.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/RunEndEncoded.java @@ -42,7 +42,7 @@ */ @SuppressWarnings("unused") public final class RunEndEncoded extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static RunEndEncoded getRootAsRunEndEncoded(ByteBuffer _bb) { return getRootAsRunEndEncoded(_bb, new RunEndEncoded()); } public static RunEndEncoded getRootAsRunEndEncoded(ByteBuffer _bb, RunEndEncoded obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Schema.java b/format/src/main/java/org/apache/arrow/flatbuf/Schema.java index 373071f4bb..0fd03b5a67 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Schema.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Schema.java @@ -39,7 +39,7 @@ */ @SuppressWarnings("unused") public final class Schema extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Schema getRootAsSchema(ByteBuffer _bb) { return getRootAsSchema(_bb, new Schema()); } public static Schema getRootAsSchema(ByteBuffer _bb, Schema obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/SparseMatrixIndexCSX.java b/format/src/main/java/org/apache/arrow/flatbuf/SparseMatrixIndexCSX.java index e5c1dc8a6f..b2c3e6d800 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/SparseMatrixIndexCSX.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/SparseMatrixIndexCSX.java @@ -38,7 +38,7 @@ */ @SuppressWarnings("unused") public final class SparseMatrixIndexCSX extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static SparseMatrixIndexCSX getRootAsSparseMatrixIndexCSX(ByteBuffer _bb) { return getRootAsSparseMatrixIndexCSX(_bb, new SparseMatrixIndexCSX()); } public static SparseMatrixIndexCSX getRootAsSparseMatrixIndexCSX(ByteBuffer _bb, SparseMatrixIndexCSX obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensor.java b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensor.java index 6ca7f29f12..4cf10e03a2 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensor.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensor.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class SparseTensor extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static SparseTensor getRootAsSparseTensor(ByteBuffer _bb) { return getRootAsSparseTensor(_bb, new SparseTensor()); } public static SparseTensor getRootAsSparseTensor(ByteBuffer _bb, SparseTensor obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCOO.java b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCOO.java index b43915c6af..e2321ae412 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCOO.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCOO.java @@ -69,7 +69,7 @@ */ @SuppressWarnings("unused") public final class SparseTensorIndexCOO extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static SparseTensorIndexCOO getRootAsSparseTensorIndexCOO(ByteBuffer _bb) { return getRootAsSparseTensorIndexCOO(_bb, new SparseTensorIndexCOO()); } public static SparseTensorIndexCOO getRootAsSparseTensorIndexCOO(ByteBuffer _bb, SparseTensorIndexCOO obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCSF.java b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCSF.java index a32e29a6e8..f2d42eb656 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCSF.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/SparseTensorIndexCSF.java @@ -38,7 +38,7 @@ */ @SuppressWarnings("unused") public final class SparseTensorIndexCSF extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static SparseTensorIndexCSF getRootAsSparseTensorIndexCSF(ByteBuffer _bb) { return getRootAsSparseTensorIndexCSF(_bb, new SparseTensorIndexCSF()); } public static SparseTensorIndexCSF getRootAsSparseTensorIndexCSF(ByteBuffer _bb, SparseTensorIndexCSF obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Struct_.java b/format/src/main/java/org/apache/arrow/flatbuf/Struct_.java index 9065d30067..4eb4bbcc7b 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Struct_.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Struct_.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class Struct_ extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Struct_ getRootAsStruct_(ByteBuffer _bb) { return getRootAsStruct_(_bb, new Struct_()); } public static Struct_ getRootAsStruct_(ByteBuffer _bb, Struct_ obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Tensor.java b/format/src/main/java/org/apache/arrow/flatbuf/Tensor.java index 7bba1e3455..4bcf93d9a0 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Tensor.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Tensor.java @@ -35,7 +35,7 @@ @SuppressWarnings("unused") public final class Tensor extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Tensor getRootAsTensor(ByteBuffer _bb) { return getRootAsTensor(_bb, new Tensor()); } public static Tensor getRootAsTensor(ByteBuffer _bb, Tensor obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/TensorDim.java b/format/src/main/java/org/apache/arrow/flatbuf/TensorDim.java index 566ec2e3a2..36459e742a 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/TensorDim.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/TensorDim.java @@ -40,7 +40,7 @@ */ @SuppressWarnings("unused") public final class TensorDim extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static TensorDim getRootAsTensorDim(ByteBuffer _bb) { return getRootAsTensorDim(_bb, new TensorDim()); } public static TensorDim getRootAsTensorDim(ByteBuffer _bb, TensorDim obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Time.java b/format/src/main/java/org/apache/arrow/flatbuf/Time.java index 0c635ea103..2d692c1006 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Time.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Time.java @@ -51,7 +51,7 @@ */ @SuppressWarnings("unused") public final class Time extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Time getRootAsTime(ByteBuffer _bb) { return getRootAsTime(_bb, new Time()); } public static Time getRootAsTime(ByteBuffer _bb, Time obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Timestamp.java b/format/src/main/java/org/apache/arrow/flatbuf/Timestamp.java index 2f33b2ee5e..df0e79fced 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Timestamp.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Timestamp.java @@ -142,7 +142,7 @@ */ @SuppressWarnings("unused") public final class Timestamp extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Timestamp getRootAsTimestamp(ByteBuffer _bb) { return getRootAsTimestamp(_bb, new Timestamp()); } public static Timestamp getRootAsTimestamp(ByteBuffer _bb, Timestamp obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Union.java b/format/src/main/java/org/apache/arrow/flatbuf/Union.java index 9c3c6b8e3e..21f3b47437 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Union.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Union.java @@ -41,7 +41,7 @@ */ @SuppressWarnings("unused") public final class Union extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Union getRootAsUnion(ByteBuffer _bb) { return getRootAsUnion(_bb, new Union()); } public static Union getRootAsUnion(ByteBuffer _bb, Union obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Utf8.java b/format/src/main/java/org/apache/arrow/flatbuf/Utf8.java index 6aa4baa7b7..1abd8ae6ce 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Utf8.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Utf8.java @@ -38,7 +38,7 @@ */ @SuppressWarnings("unused") public final class Utf8 extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Utf8 getRootAsUtf8(ByteBuffer _bb) { return getRootAsUtf8(_bb, new Utf8()); } public static Utf8 getRootAsUtf8(ByteBuffer _bb, Utf8 obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } diff --git a/format/src/main/java/org/apache/arrow/flatbuf/Utf8View.java b/format/src/main/java/org/apache/arrow/flatbuf/Utf8View.java index 1da9ef27f3..68b99a32fb 100644 --- a/format/src/main/java/org/apache/arrow/flatbuf/Utf8View.java +++ b/format/src/main/java/org/apache/arrow/flatbuf/Utf8View.java @@ -44,7 +44,7 @@ */ @SuppressWarnings("unused") public final class Utf8View extends Table { - public static void ValidateVersion() { Constants.FLATBUFFERS_24_3_25(); } + public static void ValidateVersion() { Constants.FLATBUFFERS_25_2_10(); } public static Utf8View getRootAsUtf8View(ByteBuffer _bb) { return getRootAsUtf8View(_bb, new Utf8View()); } public static Utf8View getRootAsUtf8View(ByteBuffer _bb, Utf8View obj) { _bb.order(ByteOrder.LITTLE_ENDIAN); return (obj.__assign(_bb.getInt(_bb.position()) + _bb.position(), _bb)); } public void __init(int _i, ByteBuffer _bb) { __reset(_i, _bb); } 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/expression_registry_helper.cc b/gandiva/src/main/cpp/expression_registry_helper.cc index 66b97c8b9e..21077ff1db 100644 --- a/gandiva/src/main/cpp/expression_registry_helper.cc +++ b/gandiva/src/main/cpp/expression_registry_helper.cc @@ -138,7 +138,7 @@ void ArrowToProtobuf(DataTypePtr type, gandiva::types::ExtGandivaType* gandiva_d default: // un-supported types. test ensures that // when one of these are added build breaks. - DCHECK(false); + ARROW_DCHECK(false); } } diff --git a/gandiva/src/main/cpp/jni_common.cc b/gandiva/src/main/cpp/jni_common.cc index ec1bb76234..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: @@ -751,7 +751,7 @@ Status JavaResizableBuffer::Resize(const int64_t new_size, bool shrink_to_fit) { } RETURN_NOT_OK(Reserve(new_size)); - DCHECK_GE(capacity_, new_size); + ARROW_DCHECK_GE(capacity_, new_size); size_ = new_size; return Status::OK(); } diff --git a/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java b/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java index 5c16c46e5e..9c5b22d659 100644 --- a/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java +++ b/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java @@ -188,7 +188,7 @@ public static Projector make( * @param configurationId Custom configuration created through config builder. * @return A native evaluator object that can be used to invoke these projections on a RecordBatch */ - public static Projector make( + public static synchronized Projector make( Schema schema, List exprs, SelectionVectorType selectionVectorType, @@ -314,7 +314,7 @@ public void evaluate( outColumns); } - private void evaluate( + private synchronized void evaluate( int numRows, List buffers, List buffersLayout, diff --git a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterProjectTest.java b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterProjectTest.java index 75169a37a9..80427de0f0 100644 --- a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterProjectTest.java +++ b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterProjectTest.java @@ -34,13 +34,11 @@ import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class FilterProjectTest extends BaseEvaluatorTest { @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleSV16() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); diff --git a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterTest.java b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterTest.java index a98a7cb6b5..7563465f37 100644 --- a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterTest.java +++ b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/FilterTest.java @@ -34,7 +34,6 @@ 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.Disabled; import org.junit.jupiter.api.Test; public class FilterTest extends BaseEvaluatorTest { @@ -73,7 +72,6 @@ List stringBufs(String[] strings) { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleInString() throws GandivaException, Exception { Field c1 = Field.nullable("c1", new ArrowType.Utf8()); TreeNode l1 = TreeBuilder.makeLiteral(1L); @@ -137,7 +135,6 @@ public void testSimpleInString() throws GandivaException, Exception { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleInInt() throws GandivaException, Exception { Field c1 = Field.nullable("c1", int32); @@ -181,7 +178,6 @@ public void testSimpleInInt() throws GandivaException, Exception { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleSV16() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); @@ -203,7 +199,6 @@ public void testSimpleSV16() throws GandivaException, Exception { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleSV16_AllMatched() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); @@ -233,7 +228,6 @@ public void testSimpleSV16_AllMatched() throws GandivaException, Exception { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleSV16_GreaterThan64Recs() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); @@ -265,7 +259,6 @@ public void testSimpleSV16_GreaterThan64Recs() throws GandivaException, Exceptio } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleSV32() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); @@ -287,7 +280,6 @@ public void testSimpleSV32() throws GandivaException, Exception { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testSimpleFilterWithNoOptimisation() throws GandivaException, Exception { Field a = Field.nullable("a", int32); Field b = Field.nullable("b", int32); diff --git a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorDecimalTest.java b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorDecimalTest.java index 74180c0f35..3916051224 100644 --- a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorDecimalTest.java +++ b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorDecimalTest.java @@ -42,13 +42,11 @@ import org.apache.arrow.vector.types.pojo.ArrowType.Decimal; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; public class ProjectorDecimalTest extends org.apache.arrow.gandiva.evaluator.BaseEvaluatorTest { @Test - @Disabled("GH-43576 - Fix and enable this test") public void test_add() throws GandivaException { int precision = 38; int scale = 8; @@ -116,7 +114,6 @@ public void test_add() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void test_add_literal() throws GandivaException { int precision = 2; int scale = 0; @@ -178,7 +175,6 @@ public void test_add_literal() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void test_multiply() throws GandivaException { int precision = 38; int scale = 8; @@ -248,7 +244,6 @@ public void test_multiply() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCompare() throws GandivaException { Decimal aType = new Decimal(38, 3, 128); Decimal bType = new Decimal(38, 2, 128); @@ -343,7 +338,6 @@ public void testCompare() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testRound() throws GandivaException { Decimal aType = new Decimal(38, 2, 128); Decimal aWithScaleZero = new Decimal(38, 0, 128); @@ -486,7 +480,6 @@ public void testRound() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCastToDecimal() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Decimal decimalWithScaleOne = new Decimal(38, 1, 128); @@ -613,7 +606,6 @@ public void testCastToDecimal() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCastToLong() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); @@ -666,7 +658,6 @@ public void testCastToLong() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCastToDouble() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); @@ -721,7 +712,6 @@ public void testCastToDouble() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCastToString() throws GandivaException { Decimal decimalType = new Decimal(38, 2, 128); Field dec = Field.nullable("dec", decimalType); @@ -783,7 +773,6 @@ public void testCastToString() throws GandivaException { } @Test - @Disabled("GH-43576 - Fix and enable this test") public void testCastStringToDecimal() throws GandivaException { Decimal decimalType = new Decimal(4, 2, 128); Field dec = Field.nullable("dec", decimalType); diff --git a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java index 0d86bd9e72..f2590226b1 100644 --- a/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java +++ b/gandiva/src/test/java/org/apache/arrow/gandiva/evaluator/ProjectorTest.java @@ -62,7 +62,6 @@ import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; -@Disabled("Disabled until GH-43981 is solved") public class ProjectorTest extends BaseEvaluatorTest { private Charset utf8Charset = Charset.forName("UTF-8"); 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/AllocationManager.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java index e9dd8cb9d2..22f9202008 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java @@ -74,8 +74,7 @@ protected AllocationManager(BufferAllocator accountingAllocator) { this.owningLedger = associate(accountingAllocator, false); } - @Nullable - BufferLedger getOwningLedger() { + @Nullable BufferLedger getOwningLedger() { return owningLedger; } 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/BaseAllocator.java b/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java index 20a89d0b7b..72946e7fb7 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java @@ -272,7 +272,7 @@ public ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) { final AllocationManager manager = new ForeignAllocationManager(this, allocation); final BufferLedger ledger = manager.associate(this); final ArrowBuf buf = - new ArrowBuf(ledger, /*bufferManager=*/ null, size, allocation.memoryAddress()); + new ArrowBuf(ledger, /* bufferManager= */ null, size, allocation.memoryAddress()); buf.writerIndex(size); listener.onAllocation(size); return buf; 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 a4db99f619..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. @@ -157,8 +159,7 @@ BufferAllocator newChildAllocator( * * @return parent allocator */ - @Nullable - BufferAllocator getParentAllocator(); + @Nullable BufferAllocator getParentAllocator(); /** * Returns the set of child allocators. 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 acf77547fb..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; @@ -32,6 +33,7 @@ public class MemoryUtil { org.slf4j.LoggerFactory.getLogger(MemoryUtil.class); private static final @Nullable Constructor DIRECT_BUFFER_CONSTRUCTOR; + /** The unsafe object from which to access the off-heap memory. */ private static final Unsafe UNSAFE; @@ -80,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; @@ -108,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; } } }); @@ -155,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 3796fb94bc..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/main/java/org/apache/arrow/util/Preconditions.java b/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java index 71e622b45f..16d763afba 100644 --- a/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java +++ b/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java @@ -1,18 +1,15 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at + * Copyright (C) 2007 The Guava Authors * - * http://www.apache.org/licenses/LICENSE-2.0 + * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. You may obtain a copy of the License at * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under the License + * is distributed on an "AS IS" BASIS, WITHOUT 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; 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/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarkConstants.java b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarkConstants.java new file mode 100644 index 0000000000..d4ef49d902 --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarkConstants.java @@ -0,0 +1,22 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; + +public class VariableWidthVectorBenchmarkConstants { + public static final String SHORT_VALUE = "InlineValue"; + public static final String LONG_VALUE = VariableWidthVectorBenchmarks.class.getName(); +} diff --git a/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarks.java index 0bce6569d2..17d3ac1fee 100644 --- a/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarks.java +++ b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorBenchmarks.java @@ -23,8 +23,10 @@ import org.apache.arrow.vector.holders.NullableVarCharHolder; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Level; import org.openjdk.jmh.annotations.Mode; import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; import org.openjdk.jmh.annotations.Scope; import org.openjdk.jmh.annotations.Setup; import org.openjdk.jmh.annotations.State; @@ -45,15 +47,18 @@ public class VariableWidthVectorBenchmarks { private static final int ALLOCATOR_CAPACITY = 1024 * 1024; - private static byte[] bytes = VariableWidthVectorBenchmarks.class.getName().getBytes(); + private static byte[] bytes = VariableWidthVectorBenchmarkConstants.LONG_VALUE.getBytes(); private ArrowBuf arrowBuff; private BufferAllocator allocator; private VarCharVector vector; + @Param({"1", "2", "10", "40"}) + private int step; + /** Setup benchmarks. */ - @Setup + @Setup(Level.Iteration) public void prepare() { allocator = new RootAllocator(ALLOCATOR_CAPACITY); vector = new VarCharVector("vector", allocator); @@ -63,7 +68,7 @@ public void prepare() { } /** Tear down benchmarks. */ - @TearDown + @TearDown(Level.Iteration) public void tearDown() { arrowBuff.close(); vector.close(); @@ -87,7 +92,7 @@ public int getValueCapacity() { @OutputTimeUnit(TimeUnit.MILLISECONDS) public int setSafeFromArray() { for (int i = 0; i < 500; ++i) { - vector.setSafe(i * 40, bytes); + vector.setSafe(i * step, bytes); } return vector.getBufferSize(); } diff --git a/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorInlineValueBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorInlineValueBenchmarks.java new file mode 100644 index 0000000000..354d1b1479 --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/vector/VariableWidthVectorInlineValueBenchmarks.java @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.concurrent.TimeUnit; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.holders.NullableVarCharHolder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +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.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 BaseVariableWidthVector}. */ +@State(Scope.Benchmark) +public class VariableWidthVectorInlineValueBenchmarks { + // checkstyle:off: MissingJavadocMethod + + private static final int VECTOR_CAPACITY = 16 * 1024; + + private static final int VECTOR_LENGTH = 1024; + + private static final int ALLOCATOR_CAPACITY = 1024 * 1024; + + private static final byte[] bytes = VariableWidthVectorBenchmarkConstants.SHORT_VALUE.getBytes(); + private ArrowBuf arrowBuff; + + private BufferAllocator allocator; + + private VarCharVector vector; + + @Param({"1", "2", "10", "40"}) + private int step; + + /** Setup benchmarks. */ + @Setup(Level.Iteration) + public void prepare() { + allocator = new RootAllocator(ALLOCATOR_CAPACITY); + vector = new VarCharVector("vector", allocator); + vector.allocateNew(VECTOR_CAPACITY, VECTOR_LENGTH); + arrowBuff = allocator.buffer(VECTOR_LENGTH); + arrowBuff.setBytes(0, bytes, 0, bytes.length); + } + + /** Tear down benchmarks. */ + @TearDown(Level.Iteration) + public void tearDown() { + arrowBuff.close(); + vector.close(); + allocator.close(); + } + + /** + * Test {@link BaseVariableWidthVector#getValueCapacity()}. + * + * @return useless. To avoid DCE by JIT. + */ + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.NANOSECONDS) + public int getValueCapacity() { + return vector.getValueCapacity(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromArray() { + for (int i = 0; i < 500; ++i) { + vector.setSafe(i * step, bytes); + } + return vector.getBufferSize(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromNullableVarcharHolder() { + NullableVarCharHolder nvch = new NullableVarCharHolder(); + nvch.buffer = arrowBuff; + nvch.start = 0; + nvch.end = bytes.length; + for (int i = 0; i < 50; ++i) { + nvch.isSet = 0; + for (int j = 0; j < 9; ++j) { + int idx = 10 * i + j; + vector.setSafe(idx, nvch); + } + nvch.isSet = 1; + vector.setSafe(10 * (i + 1), nvch); + } + return vector.getBufferSize(); + } + + public static void main(String[] args) throws RunnerException { + Options opt = + new OptionsBuilder() + .include(VariableWidthVectorInlineValueBenchmarks.class.getSimpleName()) + .forks(1) + .build(); + + new Runner(opt).run(); + } + // checkstyle:on: MissingJavadocMethod +} diff --git a/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorBenchmarks.java new file mode 100644 index 0000000000..b85377dcd3 --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorBenchmarks.java @@ -0,0 +1,130 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.concurrent.TimeUnit; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.holders.NullableViewVarCharHolder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +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.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 BaseVariableWidthVector}. */ +@State(Scope.Benchmark) +public class VariableWidthViewVectorBenchmarks { + // checkstyle:off: MissingJavadocMethod + + private static final int VECTOR_CAPACITY = 16 * 1024; + + private static final int VECTOR_LENGTH = 1024; + + private static final int ALLOCATOR_CAPACITY = 1024 * 1024; + + private static byte[] bytes = VariableWidthVectorBenchmarkConstants.LONG_VALUE.getBytes(); + private ArrowBuf arrowBuff; + + private BufferAllocator allocator; + + private ViewVarCharVector vector; + + @Param({"1", "2", "10", "40"}) + private int step; + + /** Setup benchmarks. */ + @Setup(Level.Iteration) + public void prepare() { + allocator = new RootAllocator(); + vector = new ViewVarCharVector("vector", allocator); + vector.allocateNew(VECTOR_CAPACITY, VECTOR_LENGTH); + arrowBuff = allocator.buffer(VECTOR_LENGTH); + arrowBuff.setBytes(0, bytes, 0, bytes.length); + } + + /** Tear down benchmarks. */ + @TearDown(Level.Iteration) + public void tearDown() { + arrowBuff.close(); + vector.close(); + allocator.close(); + } + + /** + * Test {@link BaseVariableWidthVector#getValueCapacity()}. + * + * @return useless. To avoid DCE by JIT. + */ + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.NANOSECONDS) + public int getValueCapacity() { + return vector.getValueCapacity(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromArray() { + for (int i = 0; i < 500; ++i) { + vector.setSafe(i * step, bytes); + } + return vector.getBufferSize(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromNullableVarcharHolder() { + NullableViewVarCharHolder nvch = new NullableViewVarCharHolder(); + nvch.buffer = arrowBuff; + nvch.start = 0; + nvch.end = bytes.length; + for (int i = 0; i < 50; ++i) { + nvch.isSet = 0; + for (int j = 0; j < 9; ++j) { + int idx = 10 * i + j; + vector.setSafe(idx, nvch); + } + nvch.isSet = 1; + vector.setSafe(10 * (i + 1), nvch); + } + return vector.getBufferSize(); + } + + public static void main(String[] args) throws RunnerException { + Options opt = + new OptionsBuilder() + .include(VariableWidthViewVectorBenchmarks.class.getSimpleName()) + .forks(1) + .build(); + + new Runner(opt).run(); + } + // checkstyle:on: MissingJavadocMethod +} diff --git a/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorInlineValueBenchmarks.java b/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorInlineValueBenchmarks.java new file mode 100644 index 0000000000..32143d4773 --- /dev/null +++ b/performance/src/main/java/org/apache/arrow/vector/VariableWidthViewVectorInlineValueBenchmarks.java @@ -0,0 +1,131 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.vector; + +import java.util.concurrent.TimeUnit; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.holders.NullableViewVarCharHolder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +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.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 BaseVariableWidthVector}. */ +@State(Scope.Benchmark) +public class VariableWidthViewVectorInlineValueBenchmarks { + // checkstyle:off: MissingJavadocMethod + + private static final int VECTOR_CAPACITY = 16 * 1024; + + private static final int VECTOR_LENGTH = 1024; + + private static final int ALLOCATOR_CAPACITY = 1024 * 1024; + + private static byte[] bytes = VariableWidthVectorBenchmarkConstants.SHORT_VALUE.getBytes(); + private ArrowBuf arrowBuff; + + private BufferAllocator allocator; + + private ViewVarCharVector vector; + + @Param({"1", "2", "10", "40"}) + private int step; + + /** Setup benchmarks. */ + @Setup(Level.Invocation) + public void prepare() { + allocator = new RootAllocator(); + vector = new ViewVarCharVector("vector", allocator); + vector.allocateNew(VECTOR_CAPACITY, VECTOR_LENGTH); + vector.zeroVector(); + arrowBuff = allocator.buffer(VECTOR_LENGTH); + arrowBuff.setBytes(0, bytes, 0, bytes.length); + } + + /** Tear down benchmarks. */ + @TearDown(Level.Invocation) + public void tearDown() { + arrowBuff.close(); + vector.close(); + allocator.close(); + } + + /** + * Test {@link BaseVariableWidthVector#getValueCapacity()}. + * + * @return useless. To avoid DCE by JIT. + */ + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.NANOSECONDS) + public int getValueCapacity() { + return vector.getValueCapacity(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromArray() { + for (int i = 0; i < 500; ++i) { + vector.setSafe(i * step, bytes); + } + return vector.getBufferSize(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + @OutputTimeUnit(TimeUnit.MILLISECONDS) + public int setSafeFromNullableVarcharHolder() { + NullableViewVarCharHolder nvch = new NullableViewVarCharHolder(); + nvch.buffer = arrowBuff; + nvch.start = 0; + nvch.end = bytes.length; + for (int i = 0; i < 50; ++i) { + nvch.isSet = 0; + for (int j = 0; j < 9; ++j) { + int idx = 10 * i + j; + vector.setSafe(idx, nvch); + } + nvch.isSet = 1; + vector.setSafe(10 * (i + 1), nvch); + } + return vector.getBufferSize(); + } + + public static void main(String[] args) throws RunnerException { + Options opt = + new OptionsBuilder() + .include(VariableWidthViewVectorInlineValueBenchmarks.class.getSimpleName()) + .forks(1) + .build(); + + new Runner(opt).run(); + } + // checkstyle:on: MissingJavadocMethod +} diff --git a/pom.xml b/pom.xml index f2c8d8f1f6..78137eb4e9 100644 --- a/pom.xml +++ b/pom.xml @@ -23,12 +23,12 @@ under the License. org.apache apache - 33 + 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 @@ -79,50 +80,54 @@ under the License. - scm:git:https://github.com/apache/arrow.git - scm:git:https://github.com/apache/arrow.git + scm:git:https://github.com/apache/arrow-java.git + scm:git:https://github.com/apache/arrow-java.git main - https://github.com/apache/arrow/tree/${project.scm.tag} + https://github.com/apache/arrow-java/tree/${project.scm.tag} GitHub - https://github.com/apache/arrow/issues + https://github.com/apache/arrow-java/issues + 1773644827 ${project.build.directory}/generated-sources 1.9.0 - 5.11.3 - 2.0.16 - 33.3.1-jre - 4.1.115.Final - 1.65.0 - 3.25.5 - 2.18.1 - 3.4.1 - 24.3.25 - 1.12.0 - + 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.1 + 1.17.1 + 5.23.0 + 2 - 10.20.1 + 10.23.0 true - 2.31.0 - 5.14.2 - 3.48.2 - 1.5.12 + 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.26.3 + 3.27.7 test org.immutables value-annotations - 2.10.1 + 2.12.2 provided @@ -221,6 +226,13 @@ under the License. pom import + + org.mockito + mockito-bom + ${dep.mockito-bom.version} + pom + import + ch.qos.logback logback-classic @@ -275,12 +287,6 @@ under the License. ${dep.junit.jupiter.version} test - - org.mockito - mockito-junit-jupiter - 5.14.2 - test - ch.qos.logback logback-classic @@ -308,7 +314,7 @@ under the License. org.immutables value - 2.10.1 + 2.12.2 @@ -326,8 +332,8 @@ under the License. true UTC - 1048576 + which in turn can cause OOM. Using 2MB - 1byte to simulate the defaul limit of 2^31 - 1 bytes. --> + 2097151 false @@ -346,7 +352,7 @@ under the License. org.jacoco jacoco-maven-plugin - 0.8.12 + 0.8.15 error-prone + [17,) !m2e.version @@ -933,6 +874,7 @@ under the License. -J--add-exports=jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED -J--add-opens=jdk.compiler/com.sun.tools.javac.code=ALL-UNNAMED -J--add-opens=jdk.compiler/com.sun.tools.javac.comp=ALL-UNNAMED + --should-stop=ifError=FLOW @@ -1304,7 +1246,7 @@ under the License. arrow.test.jdk-version "JDK version used for test must be specified." ^\d{2,} - "JDK version used for test must be 11, 17, 21, ..." + "JDK version used for test must be 17, 21, ..." diff --git a/tools/pom.xml b/tools/pom.xml index f06ded294a..b4d64fd435 100644 --- a/tools/pom.xml +++ b/tools/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-tools Arrow Tools @@ -54,7 +54,7 @@ under the License. commons-cli commons-cli - 1.9.0 + 1.11.0 ch.qos.logback @@ -96,6 +96,7 @@ under the License. + org.apache.maven.plugins maven-shade-plugin @@ -114,9 +115,26 @@ under the License. **/module-info.class + + *:* + + LICENSE.txt + NOTICE.txt + META-INF/*LICENSE* + META-INF/*NOTICE* + + - + + + META-INF/LICENSE.txt + src/shade/LICENSE.txt + + + META-INF/NOTICE.txt + src/shade/NOTICE.txt + diff --git a/tools/src/shade/LICENSE.txt b/tools/src/shade/LICENSE.txt new file mode 100644 index 0000000000..eef614e7f2 --- /dev/null +++ b/tools/src/shade/LICENSE.txt @@ -0,0 +1,331 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-------------------------------------------------------------------------------- + +This binary artifact contains Netty 4.1.119.Final. + +Copyright: Copyright 2014 The Netty Project +Home page: https://netty.io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Jackson 2.18.3. + +Home page: https://github.com/FasterXML/jackson +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons Codec 1.18.0. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-codec/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons Compress 1.27.1. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-compress/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons IO 2.16.1. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-io/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons Lang 3.16.0. + +Copyright: Copyright 2001-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-lang/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Luben Zstd 1.5.7-2. + +Copyright: Copyright (c) 2015-present, Luben Karavelov/ All rights reserved. +Home page: https://github.com/luben/zstd-jni +License: BSD License +License text: + +| Zstd-jni: JNI bindings to Zstd Library +| +| Copyright (c) 2015-present, Luben Karavelov/ All rights reserved. +| +| BSD License +| +| Redistribution and use in source and binary forms, with or without modification, +| are permitted provided that the following conditions are met: +| +| * Redistributions of source code must retain the above copyright notice, this +| list of conditions and the following disclaimer. +| +| * Redistributions in binary form must reproduce the above copyright notice, this +| list of conditions and the following disclaimer in the documentation and/or +| other materials provided with the distribution. +| +| THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +| ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +| WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +| DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +| ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +| (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +| LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +| ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +| (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +| SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +-------------------------------------------------------------------------------- + +This binary artifact contains Apache Commons CLI 1.9.0. + +Copyright: Copyright 2002-2024 The Apache Software Foundation +Home page: https://commons.apache.org/proper/commons-cli/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Google Flatbuffers 25.2.10. + +Home page: https://flatbuffers.dev/ +License: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- + +This binary artifact contains Slf4j 2.0.17. + +Copyright: Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +Home page: http://www.slf4j.org/ +License: MIT +License text: + +| Copyright (c) 2004-2022 QOS.ch Sarl (Switzerland) +| All rights reserved. +| +| Permission is hereby granted, free of charge, to any person obtaining +| a copy of this software and associated documentation files (the +| "Software"), to deal in the Software without restriction, including +| without limitation the rights to use, copy, modify, merge, publish, +| distribute, sublicense, and/or sell copies of the Software, and to +| permit persons to whom the Software is furnished to do so, subject to +| the following conditions: +| +| The above copyright notice and this permission notice shall be +| included in all copies or substantial portions of the Software. +| +| THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +| EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +| MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +| NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +| LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +| OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +| WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/tools/src/shade/NOTICE.txt b/tools/src/shade/NOTICE.txt new file mode 100644 index 0000000000..1f709a5336 --- /dev/null +++ b/tools/src/shade/NOTICE.txt @@ -0,0 +1,273 @@ +Apache Arrow Java +Copyright 2016-2025 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +--------------------------------------------------- + +This product includes Netty 4.1.119.Final, with the following in its NOTICE: + +| The Netty Project +| ================= +| +| Please visit the Netty web site for more information: +| +| * https://netty.io/ +| +| Copyright 2014 The Netty Project +| +| The Netty Project licenses this file to you under the Apache License, +| version 2.0 (the "License"); you may not use this file except in compliance +| with the License. You may obtain a copy of the License at: +| +| https://www.apache.org/licenses/LICENSE-2.0 +| +| Unless required by applicable law or agreed to in writing, software +| distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +| WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +| License for the specific language governing permissions and limitations +| under the License. +| +| Also, please refer to each LICENSE..txt file, which is located in +| the 'license' directory of the distribution file, for the license terms of the +| components that this product depends on. +| +| ------------------------------------------------------------------------------- +| This product contains the extensions to Java Collections Framework which has +| been derived from the works by JSR-166 EG, Doug Lea, and Jason T. Greene: +| +| * LICENSE: +| * license/LICENSE.jsr166y.txt (Public Domain) +| * HOMEPAGE: +| * http://gee.cs.oswego.edu/cgi-bin/viewcvs.cgi/jsr166/ +| * http://viewvc.jboss.org/cgi-bin/viewvc.cgi/jbosscache/experimental/jsr166/ +| +| This product contains a modified version of Robert Harder's Public Domain +| Base64 Encoder and Decoder, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.base64.txt (Public Domain) +| * HOMEPAGE: +| * http://iharder.sourceforge.net/current/java/base64/ +| +| This product contains a modified portion of 'Webbit', an event based +| WebSocket and HTTP server, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.webbit.txt (BSD License) +| * HOMEPAGE: +| * https://github.com/joewalnes/webbit +| +| This product contains a modified portion of 'SLF4J', a simple logging +| facade for Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.slf4j.txt (MIT License) +| * HOMEPAGE: +| * https://www.slf4j.org/ +| +| This product contains a modified portion of 'Apache Harmony', an open source +| Java SE, which can be obtained at: +| +| * NOTICE: +| * license/NOTICE.harmony.txt +| * LICENSE: +| * license/LICENSE.harmony.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://archive.apache.org/dist/harmony/ +| +| This product contains a modified portion of 'jbzip2', a Java bzip2 compression +| and decompression library written by Matthew J. Francis. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jbzip2.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jbzip2/ +| +| This product contains a modified portion of 'libdivsufsort', a C API library to construct +| the suffix array and the Burrows-Wheeler transformed string for any input string of +| a constant-size alphabet written by Yuta Mori. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.libdivsufsort.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/y-256/libdivsufsort +| +| This product contains a modified portion of Nitsan Wakart's 'JCTools', Java Concurrency Tools for the JVM, +| which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jctools.txt (ASL2 License) +| * HOMEPAGE: +| * https://github.com/JCTools/JCTools +| +| This product optionally depends on 'JZlib', a re-implementation of zlib in +| pure Java, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jzlib.txt (BSD style License) +| * HOMEPAGE: +| * http://www.jcraft.com/jzlib/ +| +| This product optionally depends on 'Compress-LZF', a Java library for encoding and +| decoding data in LZF format, written by Tatu Saloranta. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.compress-lzf.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/ning/compress +| +| This product optionally depends on 'lz4', a LZ4 Java compression +| and decompression library written by Adrien Grand. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lz4.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jpountz/lz4-java +| +| This product optionally depends on 'lzma-java', a LZMA Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.lzma-java.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jponge/lzma-java +| +| This product optionally depends on 'zstd-jni', a zstd-jni Java compression +| and decompression library, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.zstd-jni.txt (BSD) +| * HOMEPAGE: +| * https://github.com/luben/zstd-jni +| +| This product contains a modified portion of 'jfastlz', a Java port of FastLZ compression +| and decompression library written by William Kinney. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jfastlz.txt (MIT License) +| * HOMEPAGE: +| * https://code.google.com/p/jfastlz/ +| +| This product contains a modified portion of and optionally depends on 'Protocol Buffers', Google's data +| interchange format, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.protobuf.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/protobuf +| +| This product optionally depends on 'Bouncy Castle Crypto APIs' to generate +| a temporary self-signed X.509 certificate when the JVM does not provide the +| equivalent functionality. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.bouncycastle.txt (MIT License) +| * HOMEPAGE: +| * https://www.bouncycastle.org/ +| +| This product optionally depends on 'Snappy', a compression library produced +| by Google Inc, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.snappy.txt (New BSD License) +| * HOMEPAGE: +| * https://github.com/google/snappy +| +| This product optionally depends on 'JBoss Marshalling', an alternative Java +| serialization API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.jboss-marshalling.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/jboss-remoting/jboss-marshalling +| +| This product optionally depends on 'Caliper', Google's micro- +| benchmarking framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.caliper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/google/caliper +| +| This product optionally depends on 'Apache Commons Logging', a logging +| framework, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-logging.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/logging/ +| +| This product optionally depends on 'Apache Log4J', a logging framework, which +| can be obtained at: +| +| * LICENSE: +| * license/LICENSE.log4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://logging.apache.org/log4j/ +| +| This product optionally depends on 'Aalto XML', an ultra-high performance +| non-blocking XML processor, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.aalto-xml.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://wiki.fasterxml.com/AaltoHome +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Twitter. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hpack.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/twitter/hpack +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Cory Benfield. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.hyper-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/python-hyper/hpack/ +| +| This product contains a modified version of 'HPACK', a Java implementation of +| the HTTP/2 HPACK algorithm written by Tatsuhiro Tsujikawa. It can be obtained at: +| +| * LICENSE: +| * license/LICENSE.nghttp2-hpack.txt (MIT License) +| * HOMEPAGE: +| * https://github.com/nghttp2/nghttp2/ +| +| This product contains a modified portion of 'Apache Commons Lang', a Java library +| provides utilities for the java.lang API, which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.commons-lang.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://commons.apache.org/proper/commons-lang/ +| +| +| This product contains the Maven wrapper scripts from 'Maven Wrapper', that provides an easy way to ensure a user has everything necessary to run the Maven build. +| +| * LICENSE: +| * license/LICENSE.mvn-wrapper.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/takari/maven-wrapper +| +| This product contains the dnsinfo.h header file, that provides a way to retrieve the system DNS configuration on MacOS. +| This private header is also used by Apple's open source +| mDNSResponder (https://opensource.apple.com/tarballs/mDNSResponder/). +| +| * LICENSE: +| * license/LICENSE.dnsinfo.txt (Apple Public Source License 2.0) +| * HOMEPAGE: +| * https://www.opensource.apple.com/source/configd/configd-453.19/dnsinfo/dnsinfo.h +| +| This product optionally depends on 'Brotli4j', Brotli compression and +| decompression for Java., which can be obtained at: +| +| * LICENSE: +| * license/LICENSE.brotli4j.txt (Apache License 2.0) +| * HOMEPAGE: +| * https://github.com/hyperxpro/Brotli4j diff --git a/vector/pom.xml b/vector/pom.xml index 7cd25cd43e..9b40e8820c 100644 --- a/vector/pom.xml +++ b/vector/pom.xml @@ -22,7 +22,7 @@ under the License. org.apache.arrow arrow-java-root - 19.0.0-SNAPSHOT + 20.0.0-SNAPSHOT arrow-vector Arrow Vectors @@ -60,7 +60,7 @@ under the License. commons-codec commons-codec - 1.17.1 + 1.22.0 org.apache.arrow @@ -171,6 +171,53 @@ under the License. arrow.vector.com.google.flatbuffers + + + *:* + + META-INF/LICENSE + META-INF/NOTICE + + + + + + META-INF/LICENSE.txt + src/shade/LICENSE.txt + + + META-INF/NOTICE.txt + src/shade/NOTICE.txt + + + + + + + + org.apache.maven.plugins + maven-jar-plugin + + + codegen-jar + + jar + + package + + + 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 cc2cc618d8..4b4a17d932 100644 --- a/vector/src/main/codegen/templates/AbstractFieldWriter.java +++ b/vector/src/main/codegen/templates/AbstractFieldWriter.java @@ -107,6 +107,19 @@ 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("Cannot write ExtensionHolder"); + } + @Override + public void writeExtension(Object var1) { + this.fail("Cannot write extension object"); + } + @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 /> <#assign fields = minor.fields!type.fields /> <#assign friendlyType = (minor.friendlyType!minor.boxedType!type.boxedType) /> @@ -241,6 +254,18 @@ public MapWriter map(String name, boolean keysSorted) { fail("Map"); return null; } + + @Override + public ExtensionWriter extension(String name, ArrowType arrowType) { + fail("Extension"); + return null; + } + + @Override + public ExtensionWriter extension(ArrowType arrowType) { + fail("Extension"); + return null; + } <#list vv.types as type><#list type.minor as minor> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> diff --git a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java index 06cb235f7d..2e7792fcfe 100644 --- a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java +++ b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java @@ -293,6 +293,11 @@ public MapWriter map(boolean keysSorted) { return getWriter(MinorType.MAP, new ArrowType.Map(keysSorted)); } + @Override + public ExtensionWriter extension(ArrowType arrowType) { + return getWriter(MinorType.LIST).extension(arrowType); + } + @Override public StructWriter struct(String name) { return getWriter(MinorType.STRUCT).struct(name); @@ -318,6 +323,11 @@ public MapWriter map(String name, boolean keysSorted) { return getWriter(MinorType.STRUCT).map(name, keysSorted); } + @Override + public ExtensionWriter extension(String name, ArrowType arrowType) { + return getWriter(MinorType.STRUCT).extension(name, arrowType); + } + <#list vv.types as type><#list type.minor as minor> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> 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 e952d46f1f..a4c98d7089 100644 --- a/vector/src/main/codegen/templates/BaseWriter.java +++ b/vector/src/main/codegen/templates/BaseWriter.java @@ -61,6 +61,7 @@ public interface StructWriter extends BaseWriter { void copyReaderToField(String name, FieldReader reader); StructWriter struct(String name); + ExtensionWriter extension(String name, ArrowType arrowType); ListWriter list(String name); ListWriter listView(String name); MapWriter map(String name); @@ -79,6 +80,7 @@ public interface ListWriter extends BaseWriter { ListWriter listView(); MapWriter map(); MapWriter map(boolean keysSorted); + ExtensionWriter extension(ArrowType arrowType); void copyReader(FieldReader reader); <#list vv.types as type><#list type.minor as minor> @@ -101,6 +103,36 @@ public interface MapWriter extends ListWriter { MapWriter value(); } + public interface ExtensionWriter extends BaseWriter { + + /** + * Writes a null value. + */ + void writeNull(); + + /** + * Writes value from the given extension holder. + * + * @param holder the extension holder to write + */ + void write(ExtensionHolder holder); + + /** + * Writes the given extension type value. + * + * @param value the extension type value to write + */ + void writeExtension(Object value); + + /** + * Writes the given extension type value. + * + * @param value the extension type value to write + * @param type of the extension + */ + void writeExtension(Object value, ArrowType type); + } + public interface ScalarWriter extends <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Writer, BaseWriter {} 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 c0e686f317..11d34f72c9 100644 --- a/vector/src/main/codegen/templates/PromotableWriter.java +++ b/vector/src/main/codegen/templates/PromotableWriter.java @@ -285,6 +285,9 @@ protected void setWriter(ValueVector v) { case UNION: writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory); break; + case EXTENSIONTYPE: + writer = ((ExtensionType) vector.getField().getType()).getNewFieldWriter(vector); + break; default: writer = type.getNewFieldWriter(vector); break; @@ -316,6 +319,7 @@ protected boolean requiresArrowType(MinorType type) { || type == MinorType.MAP || type == MinorType.DURATION || type == MinorType.FIXEDSIZEBINARY + || type == MinorType.EXTENSIONTYPE || (type.name().startsWith("TIMESTAMP") && type.name().endsWith("TZ")); } @@ -536,6 +540,16 @@ public void writeLargeVarChar(String value) { getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value); } + @Override + public void writeExtension(Object value, ArrowType arrowType) { + getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value, arrowType); + } + + @Override + public void write(ExtensionHolder holder) { + getWriter(MinorType.EXTENSIONTYPE, holder.type()).write(holder); + } + @Override public void allocate() { getWriter().allocate(); diff --git a/vector/src/main/codegen/templates/StructWriters.java b/vector/src/main/codegen/templates/StructWriters.java index 3e6b9fd773..413f707c70 100644 --- a/vector/src/main/codegen/templates/StructWriters.java +++ b/vector/src/main/codegen/templates/StructWriters.java @@ -83,6 +83,9 @@ public class ${mode}StructWriter extends AbstractFieldWriter { fields.put(handleCase(child.getName()), writer); break; } + case EXTENSIONTYPE: + extension(child.getName(), child.getType()); + break; case UNION: FieldType fieldType = new FieldType(addVectorAsNullable, MinorType.UNION.getType(), null, null); UnionWriter writer = new UnionWriter(container.addOrGet(child.getName(), fieldType, UnionVector.class), getNullableStructWriterFactory()); @@ -159,6 +162,29 @@ public StructWriter struct(String name) { return writer; } + @Override + public ExtensionWriter extension(String name, ArrowType arrowType) { + String finalName = handleCase(name); + FieldWriter writer = fields.get(finalName); + if(writer == null){ + int vectorCount=container.size(); + FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null); + ExtensionTypeVector vector = container.addOrGet(name, fieldType, ExtensionTypeVector.class); + writer = new PromotableWriter(vector, container, getNullableStructWriterFactory()); + if(vectorCount != container.size()) { + writer.allocate(); + } + writer.setPosition(idx()); + fields.put(finalName, writer); + } else { + if (writer instanceof PromotableWriter) { + // ensure writers are initialized + ((PromotableWriter)writer).getWriter(MinorType.EXTENSIONTYPE, arrowType); + } + } + return (ExtensionWriter) writer; + } + @Override public void close() throws Exception { clear(); 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 3962e1d073..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; @@ -201,6 +200,17 @@ public MapWriter map(String name, boolean keysSorted) { return mapWriter; } + @Override + public ExtensionWriter extension(ArrowType arrowType) { + extensionType = arrowType; + return this; + } + + @Override + public ExtensionWriter extension(String name, ArrowType arrowType) { + return writer.extension(name, arrowType); + } + <#if listName == "LargeList"> @Override public void startList() { @@ -323,6 +333,22 @@ public void writeNull() { } } + @Override + public void writeExtension(Object value) { + writer.writeExtension(value, extensionType); + writer.setPosition(writer.idx() + 1); + } + + @Override + 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> <#list type.minor as minor> <#assign name = minor.class?cap_first /> @@ -342,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); @@ -353,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); @@ -401,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 90b55cb65e..8bbf6ae0a4 100644 --- a/vector/src/main/codegen/templates/UnionMapWriter.java +++ b/vector/src/main/codegen/templates/UnionMapWriter.java @@ -231,4 +231,39 @@ public MapWriter map() { return super.map(); } } + + @Override + public ExtensionWriter extension(ArrowType type) { + switch (mode) { + case KEY: + return entryWriter.extension(MapVector.KEY_NAME, type); + case VALUE: + return entryWriter.extension(MapVector.VALUE_NAME, type); + default: + 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 e0fd0e4644..c706591966 100644 --- a/vector/src/main/codegen/templates/UnionVector.java +++ b/vector/src/main/codegen/templates/UnionVector.java @@ -23,6 +23,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseValueVector; import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.AbstractStructVector; @@ -279,7 +280,10 @@ public StructVector getStruct() { <#if minor.class?starts_with("Decimal") || is_timestamp_tz(minor.class) || minor.class == "Duration" || minor.class == "FixedSizeBinary"> public ${name}Vector get${name}Vector() { if (${uncappedName}Vector == null) { - throw new IllegalArgumentException("No ${name} present. Provide ArrowType argument to create a new vector"); + ${uncappedName}Vector = internalStruct.getChild(fieldName(MinorType.${name?upper_case}), ${name}Vector.class); + if (${uncappedName}Vector == null) { + throw new IllegalArgumentException("No ${name} present. Provide ArrowType argument to create a new vector"); + } } return ${uncappedName}Vector; } @@ -375,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); } @@ -721,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 bfe97e2770..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,6 +215,33 @@ public MapWriter asMap(ArrowType arrowType) { return getMapWriter(arrowType); } + private java.util.Map extensionWriters = new HashMap<>(); + + private ExtensionWriter getExtensionWriter(ArrowType arrowType) { + 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) { return getWriter(minorType, null); } @@ -227,6 +256,8 @@ BaseWriter getWriter(MinorType minorType, ArrowType arrowType) { return getListViewWriter(); case MAP: return getMapWriter(arrowType); + case EXTENSIONTYPE: + return getExtensionWriter(arrowType); <#list vv.types as type> <#list type.minor as minor> <#assign name = minor.class?cap_first /> @@ -460,6 +491,20 @@ public MapWriter map(String name, boolean keysSorted) { return getStructWriter().map(name, keysSorted); } + @Override + public ExtensionWriter extension(ArrowType arrowType) { + data.setType(idx(), MinorType.EXTENSIONTYPE); + getListWriter().setPosition(idx()); + return getListWriter().extension(arrowType); + } + + @Override + public ExtensionWriter extension(String name, ArrowType arrowType) { + data.setType(idx(), MinorType.EXTENSIONTYPE); + getStructWriter().setPosition(idx()); + return getStructWriter().extension(name, arrowType); + } + <#list vv.types as type><#list type.minor as minor> <#assign lowerName = minor.class?uncap_first /> <#if lowerName == "int" ><#assign lowerName = "integer" /> 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 552a896ea8..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()}. */ @@ -496,15 +506,14 @@ private void allocateBytes(final long valueBufferSize, final int valueCount) { private ArrowBuf allocateOffsetBuffer(final long size) { ArrowBuf offsetBuffer = allocator.buffer(size); offsetBuffer.readerIndex(0); - initOffsetBuffer(); + offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } /* 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 aaccec602f..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()}. */ @@ -514,16 +524,14 @@ private ArrowBuf allocateOffsetBuffer(final long size) { final int curSize = (int) size; ArrowBuf offsetBuffer = allocator.buffer(curSize); offsetBuffer.readerIndex(0); - initOffsetBuffer(); + offsetBuffer.setZero(0, offsetBuffer.capacity()); return offsetBuffer; } /* 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); } /** @@ -571,10 +579,13 @@ public void reallocDataBuffer(long desiredAllocSize) { return; } - final long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize); + final long newAllocationSize = + Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; - checkDataBufferSize(newAllocationSize); + if (newAllocationSize < desiredAllocSize) { + checkDataBufferSize(desiredAllocSize); + } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, valueBuffer, 0, valueBuffer.capacity()); @@ -670,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); @@ -853,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; - } + 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/BaseVariableWidthViewVector.java b/vector/src/main/java/org/apache/arrow/vector/BaseVariableWidthViewVector.java index 15d2182783..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); } } @@ -550,15 +542,18 @@ public void reallocViewBuffer(long desiredAllocSize) { if (desiredAllocSize == 0) { return; } - long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize); + long newAllocationSize = Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; - checkDataBufferSize(newAllocationSize); // for each set operation, we have to allocate 16 bytes // here we are adjusting the desired allocation-based allocation size // to align with the 16bytes requirement. newAllocationSize = roundUpToMultipleOf16(newAllocationSize); + if (newAllocationSize < desiredAllocSize) { + checkDataBufferSize(desiredAllocSize); + } + final ArrowBuf newBuf = allocator.buffer(newAllocationSize); newBuf.setBytes(0, viewBuffer, 0, viewBuffer.capacity()); @@ -587,10 +582,13 @@ public void reallocViewDataBuffer(long desiredAllocSize) { return; } - final long newAllocationSize = CommonUtil.nextPowerOfTwo(desiredAllocSize); + final long newAllocationSize = + Math.min(CommonUtil.nextPowerOfTwo(desiredAllocSize), MAX_BUFFER_SIZE); assert newAllocationSize >= 1; - checkDataBufferSize(newAllocationSize); + if (newAllocationSize < desiredAllocSize) { + checkDataBufferSize(desiredAllocSize); + } final ArrowBuf newBuf = allocator.buffer(newAllocationSize); dataBuffers.add(newBuf); @@ -677,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; @@ -848,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); } /** @@ -1367,11 +1310,13 @@ protected ArrowBuf allocateOrGetLastDataBuffer(int length) { protected final void setBytes(int index, byte[] value, int start, int length) { int writePosition = index * ELEMENT_SIZE; - // to clear the memory segment of view being written to - // this is helpful in case of overwriting the value - viewBuffer.setZero(writePosition, ELEMENT_SIZE); - if (length <= INLINE_SIZE) { + // Check if the memory segment has been written, and clear it if it has been set. + // It is recommended to batch initialize the viewBuffer before setBytes. + if (viewBuffer.getLong(writePosition) != 0 || viewBuffer.getLong(writePosition + 8) != 0) { + viewBuffer.setZero(writePosition, ELEMENT_SIZE); + } + // allocate inline buffer // set length viewBuffer.setInt(writePosition, length); @@ -1411,11 +1356,13 @@ protected final void setBytes(int index, byte[] value, int start, int length) { protected final void setBytes(int index, ArrowBuf valueBuf, int start, int length) { int writePosition = index * ELEMENT_SIZE; - // to clear the memory segment of view being written to - // this is helpful in case of overwriting the value - viewBuffer.setZero(writePosition, ELEMENT_SIZE); - if (length <= INLINE_SIZE) { + // Check if the memory segment has been written, and clear it if it has been set. + // It is recommended to batch initialize the viewBuffer before setBytes. + if (viewBuffer.getLong(writePosition) != 0 || viewBuffer.getLong(writePosition + 8) != 0) { + viewBuffer.setZero(writePosition, ELEMENT_SIZE); + } + // allocate inline buffer // set length viewBuffer.setInt(writePosition, length); 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/Decimal256Vector.java b/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java index 42ad741c85..f9d7e5cb9e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java +++ b/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java @@ -58,7 +58,7 @@ public final class Decimal256Vector extends BaseFixedWidthVector public Decimal256Vector(String name, BufferAllocator allocator, int precision, int scale) { this( name, - FieldType.nullable(new ArrowType.Decimal(precision, scale, /*bitWidth=*/ TYPE_WIDTH * 8)), + FieldType.nullable(new ArrowType.Decimal(precision, scale, /* bitWidth= */ TYPE_WIDTH * 8)), allocator); } @@ -567,8 +567,11 @@ private class TransferImpl implements TransferPair { public TransferImpl(String ref, BufferAllocator allocator) { to = - new Decimal256Vector( - ref, allocator, Decimal256Vector.this.precision, Decimal256Vector.this.scale); + (Decimal256Vector.this.field != null + && Decimal256Vector.this.field.getFieldType() != null) + ? new Decimal256Vector(ref, Decimal256Vector.this.field.getFieldType(), allocator) + : new Decimal256Vector( + ref, allocator, Decimal256Vector.this.precision, Decimal256Vector.this.scale); } public TransferImpl(Field field, BufferAllocator allocator) { diff --git a/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java b/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java index b4c55680b7..9bf1812cc6 100644 --- a/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java @@ -565,7 +565,10 @@ private class TransferImpl implements TransferPair { public TransferImpl(String ref, BufferAllocator allocator) { to = - new DecimalVector(ref, allocator, DecimalVector.this.precision, DecimalVector.this.scale); + (DecimalVector.this.field != null && DecimalVector.this.field.getFieldType() != null) + ? new DecimalVector(ref, DecimalVector.this.field.getFieldType(), allocator) + : new DecimalVector( + ref, allocator, DecimalVector.this.precision, DecimalVector.this.scale); } public TransferImpl(Field field, BufferAllocator allocator) { diff --git a/vector/src/main/java/org/apache/arrow/vector/FixedSizeBinaryVector.java b/vector/src/main/java/org/apache/arrow/vector/FixedSizeBinaryVector.java index 4add729358..2005036c78 100644 --- a/vector/src/main/java/org/apache/arrow/vector/FixedSizeBinaryVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/FixedSizeBinaryVector.java @@ -149,7 +149,10 @@ public void get(int index, NullableFixedSizeBinaryHolder holder) { */ @Override public byte[] getObject(int index) { - return get(index); + if (isSet(index) == 0) { + return null; + } + return get(valueBuffer, index, byteWidth); } public int getByteWidth() { 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/LargeVarBinaryVector.java b/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java index f38627b933..fe798494c9 100644 --- a/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/LargeVarBinaryVector.java @@ -16,6 +16,8 @@ */ package org.apache.arrow.vector; +import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; + import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.LargeVarBinaryReaderImpl; @@ -95,7 +97,7 @@ public MinorType getMinorType() { */ public byte[] get(int index) { assert index >= 0; - if (isSet(index) == 0) { + if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } final long startOffset = getStartOffset(index); @@ -127,7 +129,14 @@ public void read(int index, ReusableBuffer buffer) { */ @Override public byte[] getObject(int index) { - return get(index); + if (isSet(index) == 0) { + return null; + } + final long startOffset = getStartOffset(index); + final long dataLength = getEndOffset(index) - startOffset; + final byte[] result = new byte[(int) dataLength]; + valueBuffer.getBytes(startOffset, result, 0, (int) dataLength); + return result; } /** diff --git a/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java b/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java index 07a9a172f0..b7a765f310 100644 --- a/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/LargeVarCharVector.java @@ -101,7 +101,7 @@ public Types.MinorType getMinorType() { @Override public byte[] get(int index) { assert index >= 0; - if (isSet(index) == 0) { + if (NULL_CHECKING_ENABLED && isSet(index) == 0) { return null; } final long startOffset = getStartOffset(index); @@ -120,7 +120,7 @@ public byte[] get(int index) { @Override public Text getObject(int index) { assert index >= 0; - if (NULL_CHECKING_ENABLED && isSet(index) == 0) { + if (isSet(index) == 0) { return null; } 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/ValueVector.java b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java index 0a45409eb9..3a5058256c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ValueVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ValueVector.java @@ -264,7 +264,7 @@ public interface ValueVector extends Closeable, Iterable { * Get friendly type object from the vector. * * @param index index of object to get - * @return friendly type object + * @return friendly type object, null if value is unset */ Object getObject(int index); diff --git a/vector/src/main/java/org/apache/arrow/vector/VarBinaryVector.java b/vector/src/main/java/org/apache/arrow/vector/VarBinaryVector.java index 7196e9c910..ad76504f0f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VarBinaryVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/VarBinaryVector.java @@ -128,7 +128,15 @@ public void read(int index, ReusableBuffer buffer) { */ @Override public byte[] getObject(int index) { - return get(index); + if (isSet(index) == 0) { + return null; + } + + final int startOffset = getStartOffset(index); + final int dataLength = getEndOffset(index) - startOffset; + final byte[] result = new byte[dataLength]; + valueBuffer.getBytes(startOffset, result, 0, dataLength); + return result; } /** diff --git a/vector/src/main/java/org/apache/arrow/vector/VarCharVector.java b/vector/src/main/java/org/apache/arrow/vector/VarCharVector.java index c81e34558c..5ddc8b84d2 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VarCharVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/VarCharVector.java @@ -117,7 +117,7 @@ public byte[] get(int index) { @Override public Text getObject(int index) { assert index >= 0; - if (NULL_CHECKING_ENABLED && isSet(index) == 0) { + if (isSet(index) == 0) { return null; } diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java index ecd3fb9124..9b9a890346 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java @@ -122,6 +122,10 @@ private void loadBuffers( (int) (variadicBufferLayoutCount + TypeLayout.getTypeBufferCount(field.getType())); List ownBuffers = new ArrayList<>(bufferLayoutCount); for (int j = 0; j < bufferLayoutCount; j++) { + if (!buffers.hasNext()) { + throw new IllegalArgumentException( + "no more buffers for field " + field + ". Expected " + bufferLayoutCount); + } ArrowBuf nextBuf = buffers.next(); // for vectors without nulls, the buffer is empty, so there is no need to decompress it. ArrowBuf bufferToAdd = 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/ViewVarBinaryVector.java b/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java index 80d6952e00..c41854bb5f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ViewVarBinaryVector.java @@ -122,7 +122,10 @@ public void read(int index, ReusableBuffer buffer) { */ @Override public byte[] getObject(int index) { - return get(index); + if (isSet(index) == 0) { + return null; + } + return getData(index); } /** diff --git a/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java b/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java index dc474b68e3..9ce7f85ef6 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java @@ -115,7 +115,7 @@ public byte[] get(int index) { @Override public Text getObject(int index) { assert index >= 0; - if (NULL_CHECKING_ENABLED && isSet(index) == 0) { + if (isSet(index) == 0) { return null; } 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/compare/VectorVisitor.java b/vector/src/main/java/org/apache/arrow/vector/compare/VectorVisitor.java index 989c57a0c9..a95cde5275 100644 --- a/vector/src/main/java/org/apache/arrow/vector/compare/VectorVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/compare/VectorVisitor.java @@ -76,5 +76,6 @@ default OUT visit(LargeListViewVector left, IN value) { default OUT visit(RunEndEncodedVector left, IN value) { throw new UnsupportedOperationException( "VectorVisitor for LargeListViewVector is not supported."); - }; + } + ; } 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 76682c28fe..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; } @@ -769,6 +707,7 @@ public UnionVector promoteToUnion() { } protected void invalidateReader() { + fieldReader = null; reader = null; } @@ -783,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/AbstractExtensionTypeWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java new file mode 100644 index 0000000000..fccff6c21f --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/AbstractExtensionTypeWriter.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.types.pojo.Field; + +/** + * Base {@link AbstractFieldWriter} class for an {@link + * org.apache.arrow.vector.ExtensionTypeVector}. + * + * @param a specific {@link ExtensionTypeVector}. + */ +public class AbstractExtensionTypeWriter + extends AbstractFieldWriter { + protected final T vector; + + public AbstractExtensionTypeWriter(T vector) { + this.vector = vector; + } + + @Override + public Field getField() { + return this.vector.getField(); + } + + @Override + public int getValueCapacity() { + return this.vector.getValueCapacity(); + } + + @Override + public void allocate() { + this.vector.allocateNew(); + } + + @Override + public void close() { + this.vector.close(); + } + + @Override + public void clear() { + this.vector.clear(); + } + + @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/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 new file mode 100644 index 0000000000..93796aa77e --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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; +import org.apache.arrow.vector.holders.ExtensionHolder; +import org.apache.arrow.vector.types.pojo.Field; + +public class UnionExtensionWriter extends AbstractFieldWriter { + protected ExtensionTypeVector vector; + protected FieldWriter writer; + + public UnionExtensionWriter(ExtensionTypeVector vector) { + this.vector = vector; + } + + @Override + public void allocate() { + vector.allocateNew(); + } + + @Override + public void clear() { + vector.clear(); + } + + @Override + public int getValueCapacity() { + return vector.getValueCapacity(); + } + + @Override + public Field getField() { + return vector.getField(); + } + + @Override + public void close() throws Exception { + vector.close(); + } + + @Override + public void writeExtension(Object var1) { + this.writer.writeExtension(var1); + } + + @Override + public void write(ExtensionHolder holder) { + this.writer.write(holder); + } + + @Override + public void setPosition(int index) { + super.setPosition(index); + if (this.writer != null) { + 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/complex/writer/FieldWriter.java b/vector/src/main/java/org/apache/arrow/vector/complex/writer/FieldWriter.java index 949eb35d8e..51bf106685 100644 --- a/vector/src/main/java/org/apache/arrow/vector/complex/writer/FieldWriter.java +++ b/vector/src/main/java/org/apache/arrow/vector/complex/writer/FieldWriter.java @@ -16,6 +16,7 @@ */ package org.apache.arrow.vector.complex.writer; +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.ScalarWriter; @@ -25,7 +26,8 @@ * Composite of all writer types. Writers are convenience classes for incrementally adding values to * {@linkplain org.apache.arrow.vector.ValueVector}s. */ -public interface FieldWriter extends StructWriter, ListWriter, MapWriter, ScalarWriter { +public interface FieldWriter + extends StructWriter, ListWriter, MapWriter, ScalarWriter, ExtensionWriter { void allocate(); void clear(); 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/dictionary/ListSubfieldEncoder.java b/vector/src/main/java/org/apache/arrow/vector/dictionary/ListSubfieldEncoder.java index 62b9628967..f56d819885 100644 --- a/vector/src/main/java/org/apache/arrow/vector/dictionary/ListSubfieldEncoder.java +++ b/vector/src/main/java/org/apache/arrow/vector/dictionary/ListSubfieldEncoder.java @@ -60,7 +60,7 @@ private static BaseListVector cloneVector(BaseListVector vector, BufferAllocator BaseListVector cloned = (BaseListVector) fieldType.createNewSingleVector( - vector.getField().getName(), allocator, /*schemaCallBack=*/ null); + vector.getField().getName(), allocator, /* schemaCallBack= */ null); final ArrowFieldNode fieldNode = new ArrowFieldNode(vector.getValueCount(), vector.getNullCount()); diff --git a/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java b/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java index dc25bc3268..8ff152fb1c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java +++ b/vector/src/main/java/org/apache/arrow/vector/dictionary/StructSubfieldEncoder.java @@ -80,7 +80,7 @@ private static StructVector cloneVector(StructVector vector, BufferAllocator all StructVector cloned = (StructVector) fieldType.createNewSingleVector( - vector.getField().getName(), allocator, /*schemaCallback=*/ null); + vector.getField().getName(), allocator, /* schemaCallback= */ null); final ArrowFieldNode fieldNode = new ArrowFieldNode(vector.getValueCount(), vector.getNullCount()); @@ -120,7 +120,7 @@ public StructVector encode(StructVector vector, Map columnToDicti dictionary.getEncoding().getIndexType(), dictionary.getEncoding()); childrenFields.add( - new Field(childVector.getField().getName(), indexFieldType, /*children=*/ null)); + new Field(childVector.getField().getName(), indexFieldType, /* children= */ null)); } } 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 new file mode 100644 index 0000000000..4d3f767aef --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/holders/ExtensionHolder.java @@ -0,0 +1,26 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT 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.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/table/Row.java b/vector/src/main/java/org/apache/arrow/vector/table/Row.java index b89159b5ee..8d21f36627 100644 --- a/vector/src/main/java/org/apache/arrow/vector/table/Row.java +++ b/vector/src/main/java/org/apache/arrow/vector/table/Row.java @@ -118,8 +118,10 @@ public class Row implements Iterator { /** The table we're enumerating. */ protected final BaseTable table; + /** the current row number. */ protected int rowNumber = -1; + /** Indicates whether the next non-deleted row has been determined yet. */ private boolean nextRowSet; diff --git a/vector/src/main/java/org/apache/arrow/vector/types/Types.java b/vector/src/main/java/org/apache/arrow/vector/types/Types.java index e9b963b62c..17503f98c8 100644 --- a/vector/src/main/java/org/apache/arrow/vector/types/Types.java +++ b/vector/src/main/java/org/apache/arrow/vector/types/Types.java @@ -116,6 +116,7 @@ import org.apache.arrow.vector.complex.impl.UnionLargeListViewWriter; import org.apache.arrow.vector.complex.impl.UnionLargeListWriter; import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.complex.impl.UnionMapWriter; import org.apache.arrow.vector.complex.impl.UnionWriter; import org.apache.arrow.vector.complex.impl.VarBinaryWriterImpl; import org.apache.arrow.vector.complex.impl.VarCharWriterImpl; @@ -721,7 +722,7 @@ public FieldVector getNewVector( @Override public FieldWriter getNewFieldWriter(ValueVector vector) { - return new UnionListWriter((MapVector) vector); + return new UnionMapWriter((MapVector) vector); } }, TIMESTAMPSECTZ(null) { diff --git a/vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java b/vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java index b625f602ca..2fa16c66b8 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/IntObjectHashMap.java @@ -1,18 +1,16 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at + * Copyright 2014 The Netty Project * - * http://www.apache.org/licenses/LICENSE-2.0 + * The Netty Project licenses this file to you under the Apache License, version 2.0 (the + * "License"); you may not use this file except in compliance with the License. You may obtain a + * copy of the License at: * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under the License + * is distributed on an "AS IS" BASIS, WITHOUT 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; diff --git a/vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java b/vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java index 0de31b56df..362c60b5c3 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/IntObjectMap.java @@ -1,18 +1,16 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at + * Copyright 2014 The Netty Project * - * http://www.apache.org/licenses/LICENSE-2.0 + * The Netty Project licenses this file to you under the Apache License, version 2.0 (the + * "License"); you may not use this file except in compliance with the License. You may obtain a + * copy of the License at: * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under the License + * is distributed on an "AS IS" BASIS, WITHOUT 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; 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 e703571b37..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 @@ -19,16 +19,22 @@ import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import java.util.HashSet; +import java.util.List; +import org.apache.arrow.memory.ArrowBuf; 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; @@ -37,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. */ @@ -91,7 +98,6 @@ public ValueVector visit(BaseFixedWidthVector deltaVector, Void value) { deltaVector.getDataBuffer(), deltaVector.getValueCount(), targetVector.getDataBuffer()); - } else { MemoryUtil.copyMemory( deltaVector.getDataBuffer().memoryAddress(), @@ -119,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 @@ -143,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); @@ -154,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 @@ -166,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); @@ -190,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 @@ -215,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); @@ -226,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 @@ -239,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); @@ -247,8 +262,66 @@ public ValueVector visit(BaseLargeVariableWidthVector deltaVector, Void value) { } @Override - public ValueVector visit(BaseVariableWidthViewVector left, Void value) { - throw new UnsupportedOperationException("View vectors are not supported."); + public ValueVector visit(BaseVariableWidthViewVector deltaVector, Void value) { + Preconditions.checkArgument( + typeVisitor.equals(deltaVector), + "The targetVector to append must have the same type as the targetVector being appended"); + + if (deltaVector.getValueCount() == 0) { + return targetVector; // nothing to append, return + } + + int oldTargetValueCount = targetVector.getValueCount(); + int newValueCount = oldTargetValueCount + deltaVector.getValueCount(); + + // make sure there is enough capacity + while (targetVector.getValueCapacity() < newValueCount) { + // Do not call BaseVariableWidthViewVector#reAlloc() here, + // because reallocViewDataBuffer() is always unnecessary + ((BaseVariableWidthViewVector) targetVector).reallocValidityBuffer(); + ((BaseVariableWidthViewVector) targetVector).reallocViewBuffer(); + } + + // append validity buffer + BitVectorHelper.concatBits( + targetVector.getValidityBuffer(), + oldTargetValueCount, + deltaVector.getValidityBuffer(), + deltaVector.getValueCount(), + targetVector.getValidityBuffer()); + + // append data buffers + BaseVariableWidthViewVector targetViewVector = (BaseVariableWidthViewVector) targetVector; + List targetDataBuffers = targetViewVector.getDataBuffers(); + final int oldTargetDataBufferCount = targetDataBuffers.size(); + List deltaVectorDataBuffers = deltaVector.getDataBuffers(); + deltaVectorDataBuffers.forEach(buf -> buf.getReferenceManager().retain()); + targetDataBuffers.addAll(deltaVectorDataBuffers); + + // append view buffer + ArrowBuf targetViewBuffer = targetVector.getDataBuffer(); + MemoryUtil.copyMemory( + deltaVector.getDataBuffer().memoryAddress(), + targetViewBuffer.memoryAddress() + + (long) BaseVariableWidthViewVector.ELEMENT_SIZE * oldTargetValueCount, + (long) BaseVariableWidthViewVector.ELEMENT_SIZE * deltaVector.getValueCount()); + + // update view buffer + for (int i = oldTargetValueCount; i < newValueCount; i++) { + if (targetViewVector.isSet(i) > 0 + && targetViewVector.getValueLength(i) > BaseVariableWidthViewVector.INLINE_SIZE) { + long start = + (long) i * BaseVariableWidthViewVector.ELEMENT_SIZE + + BaseVariableWidthViewVector.LENGTH_WIDTH + + BaseVariableWidthViewVector.PREFIX_WIDTH; + // shift buf id + int bufferId = targetViewBuffer.getInt(start); + targetViewBuffer.setInt(start, bufferId + oldTargetDataBufferCount); + } + } + + targetVector.setValueCount(newValueCount); + return targetVector; } @Override @@ -267,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) { @@ -308,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; @@ -336,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) { @@ -363,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++) { @@ -379,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( @@ -639,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/shade/LICENSE.txt b/vector/src/shade/LICENSE.txt new file mode 100644 index 0000000000..147a99ee2f --- /dev/null +++ b/vector/src/shade/LICENSE.txt @@ -0,0 +1,209 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +-------------------------------------------------------------------------------- + +This binary artifact contains Google Flatbuffers 25.2.10. + +Home page: https://flatbuffers.dev/ +License: https://www.apache.org/licenses/LICENSE-2.0 diff --git a/vector/src/shade/NOTICE.txt b/vector/src/shade/NOTICE.txt new file mode 100644 index 0000000000..0d96a64614 --- /dev/null +++ b/vector/src/shade/NOTICE.txt @@ -0,0 +1,5 @@ +Apache Arrow Java +Copyright 2016-2025 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java index c155ab98fa..b995dc5d92 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java @@ -26,6 +26,7 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.ArrowType; +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; @@ -375,6 +376,33 @@ public void testGetTransferPairWithField() { assertSame(fromVector.getField(), toVector.getField()); } + @Test + public void testGetTransferPairWithoutField() { + final Decimal256Vector fromVector = new Decimal256Vector("decimal", allocator, 10, scale); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final Decimal256Vector toVector = (Decimal256Vector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + } + + @Test + public void testGetTransferPairWithoutFieldNonNullable() { + final FieldType decimal256NonNullableType = + new FieldType( + false, new ArrowType.Decimal(10, scale, Decimal256Vector.TYPE_WIDTH * 8), null); + final Decimal256Vector fromVector = + new Decimal256Vector("decimal", decimal256NonNullableType, allocator); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final Decimal256Vector toVector = (Decimal256Vector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + assertSame(decimal256NonNullableType, toVector.getField().getFieldType()); + } + private void verifyWritingArrowBufWithBigEndianBytes( Decimal256Vector decimalVector, ArrowBuf buf, BigDecimal[] expectedValues, int length) { decimalVector.allocateNew(); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java index d5310bad0e..85c11e8f3d 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java @@ -26,6 +26,7 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.ArrowType; +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; @@ -371,6 +372,31 @@ public void testGetTransferPairWithField() { assertSame(fromVector.getField(), toVector.getField()); } + @Test + public void testGetTransferPairWithoutField() { + final DecimalVector fromVector = new DecimalVector("decimal", allocator, 10, scale); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final DecimalVector toVector = (DecimalVector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + } + + @Test + public void testGetTransferPairWithoutFieldNonNullable() { + final FieldType decimalNonNullableType = + new FieldType(false, new ArrowType.Decimal(10, scale), null); + final DecimalVector fromVector = + new DecimalVector("decimal", decimalNonNullableType, allocator); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final DecimalVector toVector = (DecimalVector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + } + private void verifyWritingArrowBufWithBigEndianBytes( DecimalVector decimalVector, ArrowBuf buf, BigDecimal[] expectedValues, int length) { decimalVector.allocateNew(); diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java index 9cd89d57ff..9ac30730c4 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java @@ -408,7 +408,7 @@ public void testGetFieldTypeInfo() throws Exception { final FieldType fieldType = new FieldType( - false, new ArrowType.Union(UnionMode.Dense, typeIds), /*dictionary=*/ null, metadata); + false, new ArrowType.Union(UnionMode.Dense, typeIds), /* dictionary= */ null, metadata); final Field field = new Field("union", fieldType, children); MinorType minorType = MinorType.DENSEUNION; diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java index d65047efb1..0945919b91 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDictionaryVector.java @@ -1107,9 +1107,9 @@ public void testDictionaryUInt1() { new Dictionary( dictionaryVector, new DictionaryEncoding( - /*id=*/ 10L, - /*ordered=*/ false, - /*indexType=*/ new ArrowType.Int(/*bitWidth*/ 8, /*isSigned*/ false))); + /* id= */ 10L, + /* ordered= */ false, + /* indexType= */ new ArrowType.Int(/*bitWidth*/ 8, /*isSigned*/ false))); testDictionary(dictionary1, (vector, index) -> ((UInt1Vector) vector).get(index)); } } @@ -1122,9 +1122,9 @@ public void testDictionaryUInt2() { new Dictionary( dictionaryVector, new DictionaryEncoding( - /*id=*/ 20L, - /*ordered=*/ false, - /*indexType=*/ new ArrowType.Int(/*bitWidth=*/ 16, /*isSigned*/ false))); + /* id= */ 20L, + /* ordered= */ false, + /* indexType= */ new ArrowType.Int(/* bitWidth= */ 16, /*isSigned*/ false))); testDictionary(dictionary2, (vector, index) -> ((UInt2Vector) vector).get(index)); } } @@ -1137,9 +1137,9 @@ public void testDictionaryUInt4() { new Dictionary( dictionaryVector, new DictionaryEncoding( - /*id=*/ 30L, - /*ordered=*/ false, - /*indexType=*/ new ArrowType.Int(/*bitWidth=*/ 32, /*isSigned*/ false))); + /* id= */ 30L, + /* ordered= */ false, + /* indexType= */ new ArrowType.Int(/* bitWidth= */ 32, /*isSigned*/ false))); testDictionary(dictionary4, (vector, index) -> ((UInt4Vector) vector).get(index)); } } @@ -1152,9 +1152,9 @@ public void testDictionaryUInt8() { new Dictionary( dictionaryVector, new DictionaryEncoding( - /*id=*/ 40L, - /*ordered=*/ false, - /*indexType=*/ new ArrowType.Int(/*bitWidth=*/ 64, /*isSigned*/ false))); + /* id= */ 40L, + /* ordered= */ false, + /* indexType= */ new ArrowType.Int(/* bitWidth= */ 64, /*isSigned*/ false))); testDictionary(dictionary8, (vector, index) -> (int) ((UInt8Vector) vector).get(index)); } } @@ -1174,9 +1174,9 @@ public void testDictionaryUIntOverflow() { new Dictionary( dictionaryVector, new DictionaryEncoding( - /*id=*/ 10L, - /*ordered=*/ false, - /*indexType=*/ new ArrowType.Int(/*bitWidth=*/ 8, /*isSigned*/ false))); + /* id= */ 10L, + /* ordered= */ false, + /* indexType= */ new ArrowType.Int(/* bitWidth= */ 8, /*isSigned*/ false))); try (VarCharVector vector = new VarCharVector("vector", allocator)) { setVector(vector, "255"); 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 f582406de6..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; @@ -61,7 +68,7 @@ public void terminate() throws Exception { @Test public void testIntType() { - try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /*size=*/ 2, allocator)) { + try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /* size= */ 2, allocator)) { IntVector nested = (IntVector) vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())).getVector(); @@ -90,7 +97,7 @@ public void testIntType() { @Test public void testFloatTypeNullable() { - try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /*size=*/ 2, allocator)) { + try (FixedSizeListVector vector = FixedSizeListVector.empty("list", /* size= */ 2, allocator)) { Float4Vector nested = (Float4Vector) vector.addOrGetVector(FieldType.nullable(MinorType.FLOAT4.getType())).getVector(); @@ -275,7 +282,7 @@ public void testTransferEmptyVector() throws Exception { @Test public void testConsistentChildName() throws Exception { try (FixedSizeListVector listVector = - FixedSizeListVector.empty("sourceVector", /*size=*/ 2, allocator)) { + FixedSizeListVector.empty("sourceVector", /* size= */ 2, allocator)) { String emptyListStr = listVector.getField().toString(); assertTrue(emptyListStr.contains(ListVector.DATA_VECTOR_NAME)); @@ -292,7 +299,7 @@ public void testUnionFixedSizeListWriterWithNulls() throws Exception { * Read and verify */ try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -321,7 +328,7 @@ public void testUnionFixedSizeListWriterWithNulls() throws Exception { @Test public void testUnionFixedSizeListWriter() throws Exception { try (final FixedSizeListVector vector1 = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); @@ -350,7 +357,7 @@ public void testUnionFixedSizeListWriter() throws Exception { @Test public void testWriteDecimal() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -379,7 +386,7 @@ public void testWriteDecimal() throws Exception { @Test public void testDecimalIndexCheck() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -405,7 +412,7 @@ public void testWriteIllegalData() throws Exception { IllegalStateException.class, () -> { try (final FixedSizeListVector vector1 = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); @@ -430,7 +437,7 @@ public void testWriteIllegalData() throws Exception { @Test public void testSplitAndTransfer() throws Exception { try (final FixedSizeListVector vector1 = - FixedSizeListVector.empty("vector", /*size=*/ 3, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 3, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); @@ -462,7 +469,7 @@ public void testSplitAndTransfer() throws Exception { @Test public void testZeroWidthVector() { try (final FixedSizeListVector vector1 = - FixedSizeListVector.empty("vector", /*size=*/ 0, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 0, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); @@ -494,7 +501,7 @@ public void testZeroWidthVector() { @Test public void testVectorWithNulls() { try (final FixedSizeListVector vector1 = - FixedSizeListVector.empty("vector", /*size=*/ 4, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer1 = vector1.getWriter(); writer1.allocate(); @@ -527,7 +534,7 @@ public void testVectorWithNulls() { @Test public void testWriteVarCharHelpers() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 4, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -547,7 +554,7 @@ public void testWriteVarCharHelpers() throws Exception { @Test public void testWriteLargeVarCharHelpers() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 4, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -567,7 +574,7 @@ public void testWriteLargeVarCharHelpers() throws Exception { @Test public void testWriteVarBinaryHelpers() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 4, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -599,7 +606,7 @@ public void testWriteVarBinaryHelpers() throws Exception { @Test public void testWriteLargeVarBinaryHelpers() throws Exception { try (final FixedSizeListVector vector = - FixedSizeListVector.empty("vector", /*size=*/ 4, allocator)) { + FixedSizeListVector.empty("vector", /* size= */ 4, allocator)) { UnionFixedSizeListWriter writer = vector.getWriter(); writer.allocate(); @@ -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/TestIntervalMonthDayNanoVector.java b/vector/src/test/java/org/apache/arrow/vector/TestIntervalMonthDayNanoVector.java index 2b39db3cd4..b576c145fe 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestIntervalMonthDayNanoVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestIntervalMonthDayNanoVector.java @@ -47,7 +47,7 @@ public void terminate() throws Exception { @Test public void testBasics() { try (final IntervalMonthDayNanoVector vector = - new IntervalMonthDayNanoVector(/*name=*/ "", allocator)) { + new IntervalMonthDayNanoVector(/* name= */ "", allocator)) { int valueCount = 100; vector.setInitialCapacity(valueCount); vector.allocateNew(); @@ -61,12 +61,12 @@ public void testBasics() { holder.days = Integer.MIN_VALUE; holder.nanoseconds = Long.MIN_VALUE; - vector.set(0, /*months=*/ 1, /*days=*/ 2, /*nanoseconds=*/ -2); - vector.setSafe(2, /*months=*/ 1, /*days=*/ 2, /*nanoseconds=*/ -3); - vector.setSafe(/*index=*/ 4, nullableHolder); + vector.set(0, /* months= */ 1, /* days= */ 2, /* nanoseconds= */ -2); + vector.setSafe(2, /* months= */ 1, /* days= */ 2, /* nanoseconds= */ -3); + vector.setSafe(/* index= */ 4, nullableHolder); vector.set(3, holder); nullableHolder.isSet = 0; - vector.setSafe(/*index=*/ 5, nullableHolder); + vector.setSafe(/* index= */ 5, nullableHolder); vector.setValueCount(5); assertEquals("P1M2D PT-0.000000002S ", vector.getAsStringBuilder(0).toString()); 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 a4197c50b5..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,14 +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; @@ -1241,4 +1250,416 @@ public void testMakeTransferPairPreserveNullability() { assertEquals(intField, vec.getField().getChildren().get(0)); assertEquals(intField, res.getField().getChildren().get(0)); } + + @Test + public void testMapTypeReturnsSupportedMapWriter() { + try (final MapVector vector = MapVector.empty("map", allocator, false)) { + vector.allocateNew(); + FieldWriter mapWriter = MinorType.MAP.getNewFieldWriter(vector); + + mapWriter.startMap(); + mapWriter.startEntry(); + mapWriter.key().bigInt().writeBigInt(1); + mapWriter.value().integer().writeInt(11); + mapWriter.endEntry(); + mapWriter.endMap(); + + Object result = vector.getObject(0); + ArrayList resultSet = (ArrayList) result; + Map resultStruct = (Map) resultSet.get(0); + assertEquals(1L, getResultKey(resultStruct)); + 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 4ef0fbe2d9..8c8a45f588 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java @@ -26,6 +26,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.AbstractStructVector; import org.apache.arrow.vector.complex.ListVector; @@ -34,9 +35,11 @@ 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; +import org.apache.arrow.vector.types.pojo.ArrowType; 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; @@ -157,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()); } } @@ -336,6 +345,40 @@ public void testGetTransferPairWithFieldAndCallBack() { } } + @Test + public void testStructVectorWithExtensionTypes() { + 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)); + StructVector s1 = new StructVector(structField, allocator, null); + StructVector s2 = (StructVector) structField.createVector(allocator); + s1.close(); + s2.close(); + } + + @Test + public void testStructVectorTransferPairWithExtensionType() { + 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)); + + StructVector s1 = (StructVector) structField.createVector(allocator); + UuidVector uuidVector = + s1.addOrGet("uuid_child", FieldType.nullable(uuidType), UuidVector.class); + s1.setValueCount(1); + uuidVector.set(0, new UUID(1, 2)); + s1.setIndexDefined(0); + + TransferPair tp = s1.getTransferPair(structField, allocator); + final StructVector toVector = (StructVector) tp.getTo(); + assertEquals(s1.getField(), toVector.getField()); + + s1.close(); + toVector.close(); + } + private StructVector simpleStructVector(String name, BufferAllocator allocator) { final String INT_COL = "struct_int_child"; final String FLT_COL = "struct_flt_child"; diff --git a/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java b/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java index 6c05073c16..40c05f9b11 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUnionVector.java @@ -395,7 +395,10 @@ public void testGetFieldTypeInfo() throws Exception { final FieldType fieldType = new FieldType( - false, new ArrowType.Union(UnionMode.Sparse, typeIds), /*dictionary=*/ null, metadata); + false, + new ArrowType.Union(UnionMode.Sparse, typeIds), + /* dictionary= */ null, + metadata); final Field field = new Field("union", fieldType, children); MinorType minorType = MinorType.UNION; 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 3845652ad0..d91b2004c0 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestUtils.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestUtils.java @@ -16,9 +16,12 @@ */ package org.apache.arrow.vector; +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 { @@ -52,4 +55,35 @@ public static T newVector( Class c, String name, MinorType type, BufferAllocator allocator) { return c.cast(FieldType.nullable(type.getType()).createNewSingleVector(name, allocator, null)); } + + public static String generateRandomString(int length) { + Random random = new Random(); + StringBuilder sb = new StringBuilder(length); + for (int i = 0; i < length; i++) { + sb.append(random.nextInt(10)); // 0-9 + } + 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 83e470ae25..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; @@ -95,7 +100,7 @@ public void init() { private static final byte[] STR5 = "EEE5".getBytes(utf8Charset); private static final byte[] STR6 = "FFFFF6".getBytes(utf8Charset); private static final int MAX_VALUE_COUNT = - (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 7); + (int) (Integer.getInteger("arrow.vector.max_allocation_bytes", Integer.MAX_VALUE) / 9); private static final int MAX_VALUE_COUNT_8BYTE = (int) (MAX_VALUE_COUNT / 2); @AfterEach @@ -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 a4533dba3b..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; @@ -160,7 +162,7 @@ public void testDataBufferBasedAllocationInSameBuffer() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(48, 4); final int valueCount = 4; - String str4 = generateRandomString(34); + String str4 = TestUtils.generateRandomString(34); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); @@ -216,7 +218,7 @@ public void testDataBufferBasedAllocationInOtherBuffer() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(48, 4); final int valueCount = 4; - String str4 = generateRandomString(35); + String str4 = TestUtils.generateRandomString(35); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); @@ -271,7 +273,7 @@ public void testDataBufferBasedAllocationInOtherBuffer() { public void testSetSafe() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(1, 1); - byte[] str6 = generateRandomString(40).getBytes(); + byte[] str6 = TestUtils.generateRandomString(40).getBytes(); final List strings = List.of(STR0, STR1, STR2, STR3, STR4, STR5, str6); // set data to a position out of capacity index @@ -305,8 +307,8 @@ public void testMixedAllocation() { try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { viewVarCharVector.allocateNew(128, 6); final int valueCount = 6; - String str4 = generateRandomString(35); - String str6 = generateRandomString(40); + String str4 = TestUtils.generateRandomString(35); + String str6 = TestUtils.generateRandomString(40); viewVarCharVector.set(0, STR1); viewVarCharVector.set(1, STR2); viewVarCharVector.set(2, STR3); @@ -405,7 +407,7 @@ public void testSetNullableViewVarCharHolder() { setAndCheck(viewVarCharVector, i, strings.get(size - i - 1), stringHolder); } - String longString = generateRandomString(128); + String longString = TestUtils.generateRandomString(128); setAndCheck(viewVarCharVector, 6, longString.getBytes(), stringHolder); } } @@ -441,7 +443,7 @@ public void testSetNullableViewVarBinaryHolder() { setAndCheck(viewVarBinaryVector, i, strings.get(size - i - 1), holder); } - String longString = generateRandomString(128); + String longString = TestUtils.generateRandomString(128); setAndCheck(viewVarBinaryVector, 6, longString.getBytes(), holder); } } @@ -1169,7 +1171,7 @@ public void testOverwriteShortFromLongString() { vector.setValueCount(5); // overwrite index 2 with a long string - String longString = generateRandomString(128); + String longString = TestUtils.generateRandomString(128); byte[] longStringBytes = longString.getBytes(StandardCharsets.UTF_8); // since the append-only approach is used and the remaining capacity // is not enough to store the new string; a new buffer will be allocated. @@ -1373,7 +1375,7 @@ public void testOverwriteLongFromALongerLongString() { // since a new buffer is added to the dataBuffers final ArrowBuf currentDataBuf = vector.dataBuffers.get(0); final long remainingCapacity = currentDataBuf.capacity() - currentDataBuf.writerIndex(); - String longerString = generateRandomString(35); + String longerString = TestUtils.generateRandomString(35); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); assertTrue(remainingCapacity < longerStringBytes.length); @@ -1406,7 +1408,7 @@ public void testOverwriteLongFromALongerLongString() { // the remaining capacity is enough to store in the same data buffer final ArrowBuf currentDataBuf = vector.dataBuffers.get(0); final long remainingCapacity = currentDataBuf.capacity() - currentDataBuf.writerIndex(); - String longerString = generateRandomString(24); + String longerString = TestUtils.generateRandomString(24); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); assertTrue(remainingCapacity > longerStringBytes.length); @@ -1505,7 +1507,7 @@ public void testSafeOverwriteShortFromLongString() { vector.setValueCount(5); // overwrite index 2 with a long string - String longString = generateRandomString(128); + String longString = TestUtils.generateRandomString(128); byte[] longStringBytes = longString.getBytes(StandardCharsets.UTF_8); vector.setSafe(2, longStringBytes); @@ -1671,7 +1673,7 @@ public void testSafeOverwriteLongFromALongerLongString() { vector.setSafe(2, STR7); vector.setValueCount(3); - String longerString = generateRandomString(35); + String longerString = TestUtils.generateRandomString(35); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); vector.setSafe(1, longerStringBytes); @@ -1697,7 +1699,7 @@ public void testSafeOverwriteLongFromALongerLongString() { vector.setSafe(4, STR6); vector.setValueCount(5); - String longerString = generateRandomString(24); + String longerString = TestUtils.generateRandomString(24); byte[] longerStringBytes = longerString.getBytes(StandardCharsets.UTF_8); vector.setSafe(2, longerStringBytes); @@ -1869,7 +1871,7 @@ public void testCopyFromWithNulls( // to avoid re-allocation. This is to test copyFrom() without re-allocation. final int numberOfValues = initialCapacity / 2 / ViewVarCharVector.ELEMENT_SIZE; - final String prefixString = generateRandomString(12); + final String prefixString = TestUtils.generateRandomString(12); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { @@ -1965,7 +1967,7 @@ public void testCopyFromSafeWithNulls( final int numberOfValues = initialCapacity / ViewVarCharVector.ELEMENT_SIZE; - final String prefixString = generateRandomString(12); + final String prefixString = TestUtils.generateRandomString(12); for (int i = 0; i < numberOfValues; i++) { if (i % 3 == 0) { @@ -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); } } @@ -2746,7 +2748,7 @@ private void testSplitAndTransferWithMultipleDataBuffersHelper( */ @Test public void testSplitAndTransferWithMultipleDataBuffers() { - final String str4 = generateRandomString(35); + final String str4 = TestUtils.generateRandomString(35); final byte[][] data = new byte[][] {STR1, STR2, STR3, str4.getBytes(StandardCharsets.UTF_8)}; final int startIndex = 1; final int length = 3; @@ -2852,12 +2854,17 @@ public void testVectorLoadUnloadOnMixedTypes() { } } - private String generateRandomString(int length) { - Random random = new Random(); - StringBuilder sb = new StringBuilder(length); - for (int i = 0; i < length; i++) { - sb.append(random.nextInt(10)); // 0-9 + @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")); } - return sb.toString(); } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java index f5ec42c71c..bc47150376 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorReAlloc.java @@ -24,6 +24,7 @@ import java.nio.charset.StandardCharsets; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.vector.complex.DenseUnionVector; import org.apache.arrow.vector.complex.FixedSizeListVector; import org.apache.arrow.vector.complex.ListVector; @@ -222,6 +223,17 @@ public void testVariableAllocateAfterReAlloc() throws Exception { } } + @Test + public void testVariableReAllocAbove1GB() throws Exception { + try (final VarCharVector vector = new VarCharVector("", allocator)) { + long desiredSizeAboveLastPowerOf2 = + CommonUtil.nextPowerOfTwo(BaseVariableWidthVector.MAX_ALLOCATION_SIZE) / 2 + 1; + vector.reallocDataBuffer(desiredSizeAboveLastPowerOf2); + + assertTrue(vector.getDataBuffer().capacity() >= desiredSizeAboveLastPowerOf2); + } + } + @Test public void testLargeVariableAllocateAfterReAlloc() throws Exception { try (final LargeVarCharVector vector = new LargeVarCharVector("", allocator)) { 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 50f61d311e..bd3113f8bc 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java @@ -91,8 +91,8 @@ private void checkCount(BitVector vec1, IntVector vec2, VectorSchemaRoot vsr, in } private VectorSchemaRoot createBatch() { - FieldType varCharType = new FieldType(true, new ArrowType.Utf8(), /*dictionary=*/ null); - FieldType listType = new FieldType(true, new ArrowType.List(), /*dictionary=*/ null); + FieldType varCharType = new FieldType(true, new ArrowType.Utf8(), /* dictionary= */ null); + FieldType listType = new FieldType(true, new ArrowType.List(), /* dictionary= */ null); // create the schema List schemaFields = new ArrayList<>(); @@ -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/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 19b26b6d0e..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 @@ -21,15 +21,20 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import java.math.BigDecimal; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.StandardCharsets; import java.util.Objects; +import java.util.UUID; import org.apache.arrow.memory.ArrowBuf; 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; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.ListVector; @@ -37,17 +42,24 @@ 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; +import org.apache.arrow.vector.holders.NullableIntHolder; 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.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; @@ -91,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); @@ -121,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); @@ -160,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(); @@ -729,4 +758,95 @@ public void testPromoteLargeVarBinaryHelpersDirect() throws Exception { assertEquals("row4", new String(Objects.requireNonNull(uv.get(3)), StandardCharsets.UTF_8)); } } + + @Test + public void testPromoteToUnionFromDecimal() throws Exception { + try (final NonNullableStructVector container = + NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); + final DecimalVector v = + container.addOrGet( + "dec", FieldType.nullable(new ArrowType.Decimal(38, 1, 128)), DecimalVector.class); + final PromotableWriter writer = new PromotableWriter(v, container)) { + + container.allocateNew(); + container.setValueCount(1); + + writer.setPosition(0); + writer.writeDecimal(new BigDecimal("0.1")); + writer.setPosition(1); + writer.writeInt(1); + + container.setValueCount(3); + + UnionVector unionVector = (UnionVector) container.getChild("dec"); + UnionHolder holder = new UnionHolder(); + + unionVector.get(0, holder); + NullableDecimalHolder decimalHolder = new NullableDecimalHolder(); + holder.reader.read(decimalHolder); + + assertEquals(1, decimalHolder.isSet); + assertEquals( + new BigDecimal("0.1"), + DecimalUtility.getBigDecimalFromArrowBuf( + decimalHolder.buffer, 0, decimalHolder.scale, 128)); + + unionVector.get(1, holder); + NullableIntHolder intHolder = new NullableIntHolder(); + holder.reader.read(intHolder); + + assertEquals(1, intHolder.isSet); + assertEquals(1, intHolder.value); + } + } + + @Test + public void testExtensionType() throws Exception { + try (final NonNullableStructVector container = + NonNullableStructVector.empty(EMPTY_SCHEMA_PATH, allocator); + final UuidVector v = + 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.setPosition(0); + writer.writeExtension(u1, UuidType.INSTANCE); + writer.setPosition(1); + writer.writeExtension(u2, UuidType.INSTANCE); + + container.setValueCount(2); + + UuidVector uuidVector = (UuidVector) container.getChild("uuid"); + assertEquals(u1, uuidVector.getObject(0)); + 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/ipc/MessageSerializerTest.java b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java index 0a41b6c599..b529ca645a 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java @@ -208,7 +208,7 @@ public void testSerializeRecordBatchV5() throws Exception { { byte[] validBytes = out.toByteArray(); - byte[] missingBytes = Arrays.copyOfRange(validBytes, /*from=*/ 0, validBytes.length - 1); + byte[] missingBytes = Arrays.copyOfRange(validBytes, /* from= */ 0, validBytes.length - 1); ByteArrayInputStream in = new ByteArrayInputStream(missingBytes); ReadChannel channel = new ReadChannel(Channels.newChannel(in)); diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java b/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java index 74ff95d41d..dc3613df72 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/TestArrowReaderWriter.java @@ -144,19 +144,19 @@ public void init() { dictionary1 = new Dictionary( dictionaryVector1, - new DictionaryEncoding(/*id=*/ 1L, /*ordered=*/ false, /*indexType=*/ null)); + new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); dictionary2 = new Dictionary( dictionaryVector2, - new DictionaryEncoding(/*id=*/ 2L, /*ordered=*/ false, /*indexType=*/ null)); + new DictionaryEncoding(/* id= */ 2L, /* ordered= */ false, /* indexType= */ null)); dictionary3 = new Dictionary( dictionaryVector3, - new DictionaryEncoding(/*id=*/ 1L, /*ordered=*/ false, /*indexType=*/ null)); + new DictionaryEncoding(/* id= */ 1L, /* ordered= */ false, /* indexType= */ null)); dictionary4 = new Dictionary( dictionaryVector4, - new DictionaryEncoding(/*id=*/ 3L, /*ordered=*/ false, /*indexType=*/ null)); + new DictionaryEncoding(/* id= */ 3L, /* ordered= */ false, /* indexType= */ null)); } @AfterEach diff --git a/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java b/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java index f599dfa539..849fe6d667 100644 --- a/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java +++ b/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java @@ -60,6 +60,7 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VariableWidthFieldVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.BaseLargeRepeatedValueViewVector; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.BaseRepeatedValueViewVector; @@ -606,6 +607,18 @@ public static void setVector(VarCharVector vector, String... values) { vector.setValueCount(length); } + /** Populate values for ViewVarCharVector. */ + public static void setVector(ViewVarCharVector vector, String... values) { + final int length = values.length; + vector.allocateNewSafe(); + for (int i = 0; i < length; i++) { + if (values[i] != null) { + vector.setSafe(i, values[i].getBytes(StandardCharsets.UTF_8)); + } + } + vector.setValueCount(length); + } + /** Populate values for LargeVarCharVector. */ public static void setVector(LargeVarCharVector vector, String... values) { final int length = values.length; 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 8f54a6e5d7..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; @@ -41,11 +42,15 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; 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; @@ -58,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(); @@ -88,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()); @@ -112,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(); @@ -134,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())); @@ -152,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()); @@ -253,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 = @@ -295,75 +300,6 @@ public void testVectorCompare() { } } - 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 - implements ValueIterableVector { - - 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()); - } - } - static class LocationType extends ExtensionType { @Override @@ -399,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/util/TestVectorAppender.java b/vector/src/test/java/org/apache/arrow/vector/util/TestVectorAppender.java index 19eafd1b20..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 @@ -24,16 +24,25 @@ import java.util.Arrays; import java.util.Collections; 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; import org.apache.arrow.vector.Float4Vector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.LargeVarCharVector; +import org.apache.arrow.vector.TestUtils; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.compare.Range; import org.apache.arrow.vector.compare.RangeEqualsVisitor; import org.apache.arrow.vector.compare.TypeEqualsVisitor; @@ -41,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; @@ -171,6 +182,194 @@ 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; + final int length2 = 5; + try (ViewVarCharVector target = new ViewVarCharVector("", allocator); + ViewVarCharVector delta = new ViewVarCharVector("", allocator)) { + target.allocateNew(5, length1); + delta.allocateNew(5, length2); + + ValueVectorDataPopulator.setVector( + target, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9"); + ValueVectorDataPopulator.setVector(delta, "a10", "a11", "a12", "a13", null); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + try (ViewVarCharVector expected = new ViewVarCharVector("expected", allocator)) { + expected.allocateNew(); + ValueVectorDataPopulator.setVector( + expected, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9", "a10", "a11", + "a12", "a13", null); + assertVectorsEqual(expected, target); + } + } + } + + @Test + public void testAppendEmptyVariableWidthViewVector() { + try (ViewVarCharVector target = new ViewVarCharVector("", allocator); + ViewVarCharVector delta = new ViewVarCharVector("", allocator)) { + ValueVectorDataPopulator.setVector( + target, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9"); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + try (ViewVarCharVector expected = new ViewVarCharVector("expected", allocator)) { + ValueVectorDataPopulator.setVector( + expected, "a0", "a1", "a2", "a3", null, "a5", "a6", "a7", "a8", "a9"); + assertVectorsEqual(expected, target); + } + } + } + + @Test + public void testAppendShortLongVariableWidthViewVector() { + try (ViewVarCharVector target = new ViewVarCharVector("", allocator); + ViewVarCharVector delta = new ViewVarCharVector("", allocator)) { + String[] targetValues = + IntStream.range(-5, 5) + .mapToObj( + i -> TestUtils.generateRandomString(BaseVariableWidthViewVector.INLINE_SIZE + i)) + .toArray(String[]::new); + ValueVectorDataPopulator.setVector(target, targetValues); + + String[] deltaValues = + IntStream.range(-3, 3) + .mapToObj( + i -> TestUtils.generateRandomString(BaseVariableWidthViewVector.INLINE_SIZE + i)) + .toArray(String[]::new); + ValueVectorDataPopulator.setVector(delta, deltaValues); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(2, target.getDataBuffers().size()); + try (ViewVarCharVector expected = new ViewVarCharVector("expected", allocator)) { + ValueVectorDataPopulator.setVector( + expected, + Stream.concat(Arrays.stream(targetValues), Arrays.stream(deltaValues)) + .toArray(String[]::new)); + assertVectorsEqual(expected, target); + } + } + } + + @Test + public void testAppendLongVariableWidthViewVector() { + try (ViewVarCharVector target = new ViewVarCharVector("", allocator); + ViewVarCharVector delta = new ViewVarCharVector("", allocator)) { + + String[] targetValues = randomLongViewVarCharVector(target); + String[] deltaValues = randomLongViewVarCharVector(delta); + + VectorAppender appender = new VectorAppender(target); + delta.accept(appender, null); + + assertEquals(4, target.getDataBuffers().size()); + try (ViewVarCharVector expected = new ViewVarCharVector("expected", allocator)) { + ValueVectorDataPopulator.setVector( + expected, + Stream.concat(Arrays.stream(targetValues), Arrays.stream(deltaValues)) + .toArray(String[]::new)); + assertVectorsEqual(expected, target); + } + } + } + + private static String[] randomLongViewVarCharVector(ViewVarCharVector target) { + assertEquals(0, target.getDataBuffers().size()); + int initial = 64; + int stringCount = 128; + target.setInitialCapacity(initial); + String[] targetValues = + IntStream.range(0, stringCount) + .mapToObj(i -> TestUtils.generateRandomString(BaseVariableWidthViewVector.ELEMENT_SIZE)) + .toArray(String[]::new); + ValueVectorDataPopulator.setVector(target, targetValues); + assertEquals(2, target.getDataBuffers().size()); + return targetValues; + } + @Test public void testAppendEmptyVariableWidthVector() { try (VarCharVector target = new VarCharVector("", allocator); @@ -192,7 +391,15 @@ public void testAppendEmptyVariableWidthVector() { @Test public void testAppendLargeAndSmallVariableVectorsWithinLimit() { - int sixteenthOfMaxAllocation = Math.toIntExact(BaseValueVector.MAX_ALLOCATION_SIZE / 16); + // Using the max power of 2 allocation size to avoid hitting the max limit at round ups + long maxPowerOfTwoAllocationSize = + CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE); + if (maxPowerOfTwoAllocationSize > BaseValueVector.MAX_ALLOCATION_SIZE) { + maxPowerOfTwoAllocationSize = + CommonUtil.nextPowerOfTwo(BaseValueVector.MAX_ALLOCATION_SIZE / 2); + } + + int sixteenthOfMaxAllocation = Math.toIntExact(maxPowerOfTwoAllocationSize / 16); try (VarCharVector target = makeVarCharVec(1, sixteenthOfMaxAllocation); VarCharVector delta = makeVarCharVec(sixteenthOfMaxAllocation, 1)) { new VectorAppender(delta).visit(target, null); @@ -304,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); @@ -899,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;